"""
Unit tests for ReflectionNode and RewritingNode.
"""

import unittest
from unittest.mock import MagicMock, patch

from src.nodes import ReflectionNode, RewritingNode


class TestNodes(unittest.TestCase):
    @patch("src.llm_integration.get_llm")
    def test_reflection_node(self, mock_get_llm):
        # Mock LLM to return a fixed reflection
        mock_llm = MagicMock()
        mock_llm.return_value = "This is a reflection."
        mock_get_llm.return_value = mock_llm

        node = ReflectionNode("test_reflection")
        input_text = "Sample input text."
        output = node.process(input_text)

        self.assertIsInstance(output, dict)
        self.assertIn("reflection", output)
        self.assertEqual(output["reflection"], "This is a reflection.")
        # Ensure LLM was called with correct prompt
        expected_prompt = (
            "Please reflect on the following text:\n\nSample input text.\n\nReflection:"
        )
        mock_llm.assert_called_once_with(expected_prompt)

    @patch("src.llm_integration.get_llm")
    def test_rewriting_node(self, mock_get_llm):
        # Mock LLM to return a fixed rewritten text
        mock_llm = MagicMock()
        mock_llm.return_value = "Rewritten text."
        mock_get_llm.return_value = mock_llm

        node = RewritingNode("test_rewriting", style="formal")
        input_data = {"reflection": "This is a reflection."}
        output = node.process(input_data)

        self.assertIsInstance(output, dict)
        self.assertIn("rewritten", output)
        self.assertEqual(output["rewritten"], "Rewritten text.")
        # Ensure LLM was called with correct prompt
        expected_prompt = (
            "Rewrite the following reflection in a formal style:\n\nThis is a reflection.\n\nRewritten:"
        )
        mock_llm.assert_called_once_with(expected_prompt)


if __name__ == "__main__":
    unittest.main()