54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
"""
|
|
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() |