Spaces:
Sleeping
Sleeping
File size: 3,109 Bytes
f974658 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import unittest
from unittest.mock import patch, MagicMock
from graph.workflow import LangGraphWorkflow, WorkflowState
class TestLangGraphWorkflow(unittest.TestCase):
def setUp(self):
# Create mocks for agents
self.router_agent_patch = patch('graph.workflow.RouterAgent')
self.weather_agent_patch = patch('graph.workflow.WeatherAgent')
self.rag_agent_patch = patch('graph.workflow.RAGAgent')
self.evaluator_patch = patch('graph.workflow.LangSmithEvaluator')
self.mock_router_agent_class = self.router_agent_patch.start()
self.mock_weather_agent_class = self.weather_agent_patch.start()
self.mock_rag_agent_class = self.rag_agent_patch.start()
self.mock_evaluator_class = self.evaluator_patch.start()
self.mock_router_agent = self.mock_router_agent_class.return_value
self.mock_weather_agent = self.mock_weather_agent_class.return_value
self.mock_rag_agent = self.mock_rag_agent_class.return_value
self.mock_evaluator = self.mock_evaluator_class.return_value
# Initialize workflow
self.workflow = LangGraphWorkflow()
def tearDown(self):
self.router_agent_patch.stop()
self.weather_agent_patch.stop()
self.rag_agent_patch.stop()
self.evaluator_patch.stop()
def test_route_to_weather(self):
# Configure mock
self.mock_router_agent.route_query.return_value = "weather"
# Create state
state = WorkflowState(query="What's the weather in London?")
# Call the method
result = self.workflow.route(state)
# Assertions
self.assertEqual(result.action, "weather")
self.mock_router_agent.route_query.assert_called_once_with("What's the weather in London?")
def test_route_to_document(self):
# Configure mock
self.mock_router_agent.route_query.return_value = "document"
# Create state
state = WorkflowState(query="What is LangChain?")
# Call the method
result = self.workflow.route(state)
# Assertions
self.assertEqual(result.action, "document")
self.mock_router_agent.route_query.assert_called_once_with("What is LangChain?")
def test_process_weather(self):
# Configure mock
self.mock_weather_agent.get_weather_response.return_value = {
"city": "London",
"weather_data": {"temp": 15.5},
"response": "The weather in London is 15.5°C."
}
# Create state
state = WorkflowState(query="What's the weather in London?", action="weather")
# Call the method
result = self.workflow.process_weather(state)
# Assertions
self.assertEqual(result.city, "London")
self.assertEqual(result.weather_data, {"temp": 15.5})
self.assertEqual(result.response, "The weather in London is 15.5°C.") |