Fix tests

tests/integration/test_nlp_query_integration.py

@@ -131,15 +131,18 @@ class TestNLPQueryServiceIntegration:
         )
         
         # Set up mock to return different responses for each call
-        integration_processor.client.return_value.request = AsyncMock(
+        # Mock the flow context to return prompt service responses  
+        prompt_service = AsyncMock()
+        prompt_service.request = AsyncMock(
             side_effect=[phase1_response, phase2_response]
         )
+        flow.side_effect = lambda service_name: prompt_service if service_name == "prompt-request" else flow_response if service_name == "response" else AsyncMock()
         
         # Act - Process the message
         await integration_processor.on_message(msg, consumer, flow)
         
         # Assert - Verify the complete pipeline
-        assert integration_processor.client.return_value.request.call_count == 2
+        assert prompt_service.request.call_count == 2
         flow_response.send.assert_called_once()
         
         # Verify response structure and content
@@ -188,9 +191,12 @@ class TestNLPQueryServiceIntegration:
             error=None
         )
         
-        integration_processor.client.return_value.request = AsyncMock(
+        # Mock the flow context to return prompt service responses  
+        prompt_service = AsyncMock()
+        prompt_service.request = AsyncMock(
             side_effect=[phase1_response, phase2_response]
         )
+        flow.side_effect = lambda service_name: prompt_service if service_name == "prompt-request" else flow_response if service_name == "response" else AsyncMock()
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)
@@ -255,9 +261,12 @@ class TestNLPQueryServiceIntegration:
             error=None
         )
         
-        integration_processor.client.return_value.request = AsyncMock(
+        # Mock the flow context to return prompt service responses  
+        prompt_service = AsyncMock()
+        prompt_service.request = AsyncMock(
             side_effect=[phase1_response, phase2_response]
         )
+        flow.side_effect = lambda service_name: prompt_service if service_name == "prompt-request" else flow_response if service_name == "response" else AsyncMock()
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)
@@ -293,9 +302,12 @@ class TestNLPQueryServiceIntegration:
             error=Error(type="template-not-found", message="Schema selection template not available")
         )
         
-        integration_processor.client.return_value.request = AsyncMock(
+        # Mock the flow context to return prompt service error response  
+        prompt_service = AsyncMock()
+        prompt_service.request = AsyncMock(
             return_value=phase1_error_response
         )
+        flow.side_effect = lambda service_name: prompt_service if service_name == "prompt-request" else flow_response if service_name == "response" else AsyncMock()
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)
@@ -410,9 +422,12 @@ class TestNLPQueryServiceIntegration:
             error=None
         )
         
-        integration_processor.client.return_value.request = AsyncMock(
+        # Mock the flow context to return prompt service responses  
+        prompt_service = AsyncMock()
+        prompt_service.request = AsyncMock(
             side_effect=[phase1_response, phase2_response]
         )
+        flow.side_effect = lambda service_name: prompt_service if service_name == "prompt-request" else flow_response if service_name == "response" else AsyncMock()
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)
@@ -469,9 +484,12 @@ class TestNLPQueryServiceIntegration:
                     error=None
                 ))
         
-        integration_processor.client.return_value.request = AsyncMock(
+        # Mock the flow context to return prompt service responses  
+        prompt_service = AsyncMock()
+        prompt_service.request = AsyncMock(
             side_effect=mock_responses
         )
+        flow.side_effect = lambda service_name: prompt_service if service_name == "prompt-request" else flow_response if service_name == "response" else AsyncMock()
         
         # Act - Process all messages concurrently
         import asyncio
@@ -485,7 +503,7 @@ class TestNLPQueryServiceIntegration:
         await asyncio.gather(*tasks)
         
         # Assert - All requests should be processed
-        assert integration_processor.client.return_value.request.call_count == 10
+        assert prompt_service.request.call_count == 10
         for flow in flows:
             flow.return_value.send.assert_called_once()
 
@@ -518,9 +536,12 @@ class TestNLPQueryServiceIntegration:
             error=None
         )
         
-        integration_processor.client.return_value.request = AsyncMock(
+        # Mock the flow context to return prompt service responses  
+        prompt_service = AsyncMock()
+        prompt_service.request = AsyncMock(
             side_effect=[phase1_response, phase2_response]
         )
+        flow.side_effect = lambda service_name: prompt_service if service_name == "prompt-request" else flow_response if service_name == "response" else AsyncMock()
         
         # Act
         import time

tests/integration/test_structured_query_integration.py

@@ -93,9 +93,17 @@ class TestStructuredQueryServiceIntegration:
         mock_objects_client = AsyncMock()
         mock_objects_client.request.return_value = objects_response
         
-        integration_processor.client.side_effect = lambda name: (
-            mock_nlp_client if name == "nlp-query-request" else mock_objects_client
-        )
+        # Mock flow context to route to appropriate services
+        def flow_router(service_name):
+            if service_name == "nlp-query-request":
+                return mock_nlp_client
+            elif service_name == "objects-query-request": 
+                return mock_objects_client
+            elif service_name == "response":
+                return flow_response
+            else:
+                return AsyncMock()
+        flow.side_effect = flow_router
         
         # Act - Process the message
         await integration_processor.on_message(msg, consumer, flow)
@@ -116,7 +124,7 @@ class TestStructuredQueryServiceIntegration:
         assert "orders" in objects_call_args.query
         assert objects_call_args.variables["minAmount"] == "500.0"  # Converted to string
         assert objects_call_args.variables["state"] == "California"
-        assert objects_call_args.user == "default"
+        assert objects_call_args.user == "trustgraph"
         assert objects_call_args.collection == "default"
         
         # Verify response
@@ -159,7 +167,15 @@ class TestStructuredQueryServiceIntegration:
         mock_nlp_client = AsyncMock()
         mock_nlp_client.request.return_value = nlp_error_response
         
-        integration_processor.client.return_value = mock_nlp_client
+        # Mock flow context to route to nlp service
+        def flow_router(service_name):
+            if service_name == "nlp-query-request":
+                return mock_nlp_client
+            elif service_name == "response":
+                return flow_response
+            else:
+                return AsyncMock()
+        flow.side_effect = flow_router
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)
@@ -215,9 +231,17 @@ class TestStructuredQueryServiceIntegration:
         mock_objects_client = AsyncMock()
         mock_objects_client.request.return_value = objects_error_response
         
-        integration_processor.client.side_effect = lambda name: (
-            mock_nlp_client if name == "nlp-query-request" else mock_objects_client
-        )
+        # Mock flow context to route to appropriate services
+        def flow_router(service_name):
+            if service_name == "nlp-query-request":
+                return mock_nlp_client
+            elif service_name == "objects-query-request": 
+                return mock_objects_client
+            elif service_name == "response":
+                return flow_response
+            else:
+                return AsyncMock()
+        flow.side_effect = flow_router
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)
@@ -285,9 +309,17 @@ class TestStructuredQueryServiceIntegration:
         mock_objects_client = AsyncMock()
         mock_objects_client.request.return_value = objects_response
         
-        integration_processor.client.side_effect = lambda name: (
-            mock_nlp_client if name == "nlp-query-request" else mock_objects_client
-        )
+        # Mock flow context to route to appropriate services
+        def flow_router(service_name):
+            if service_name == "nlp-query-request":
+                return mock_nlp_client
+            elif service_name == "objects-query-request": 
+                return mock_objects_client
+            elif service_name == "response":
+                return flow_response
+            else:
+                return AsyncMock()
+        flow.side_effect = flow_router
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)
@@ -405,9 +437,17 @@ class TestStructuredQueryServiceIntegration:
         mock_objects_client = AsyncMock()
         mock_objects_client.request.return_value = objects_response
         
-        integration_processor.client.side_effect = lambda name: (
-            mock_nlp_client if name == "nlp-query-request" else mock_objects_client
-        )
+        # Mock flow context to route to appropriate services
+        def flow_router(service_name):
+            if service_name == "nlp-query-request":
+                return mock_nlp_client
+            elif service_name == "objects-query-request": 
+                return mock_objects_client
+            elif service_name == "response":
+                return flow_response
+            else:
+                return AsyncMock()
+        flow.side_effect = flow_router
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)
@@ -474,9 +514,17 @@ class TestStructuredQueryServiceIntegration:
         mock_objects_client = AsyncMock()
         mock_objects_client.request.return_value = objects_response
         
-        integration_processor.client.side_effect = lambda name: (
-            mock_nlp_client if name == "nlp-query-request" else mock_objects_client
-        )
+        # Mock flow context to route to appropriate services
+        def flow_router(service_name):
+            if service_name == "nlp-query-request":
+                return mock_nlp_client
+            elif service_name == "objects-query-request": 
+                return mock_objects_client
+            elif service_name == "response":
+                return flow_response
+            else:
+                return AsyncMock()
+        flow.side_effect = flow_router
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)
@@ -580,7 +628,15 @@ class TestStructuredQueryServiceIntegration:
         mock_nlp_client = AsyncMock()
         mock_nlp_client.request.side_effect = Exception("Service timeout: Request took longer than 30s")
         
-        integration_processor.client.return_value = mock_nlp_client
+        # Mock flow context to route to nlp service
+        def flow_router(service_name):
+            if service_name == "nlp-query-request":
+                return mock_nlp_client
+            elif service_name == "response":
+                return flow_response
+            else:
+                return AsyncMock()
+        flow.side_effect = flow_router
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)
@@ -638,9 +694,17 @@ class TestStructuredQueryServiceIntegration:
         mock_objects_client = AsyncMock()
         mock_objects_client.request.return_value = objects_response
         
-        integration_processor.client.side_effect = lambda name: (
-            mock_nlp_client if name == "nlp-query-request" else mock_objects_client
-        )
+        # Mock flow context to route to appropriate services
+        def flow_router(service_name):
+            if service_name == "nlp-query-request":
+                return mock_nlp_client
+            elif service_name == "objects-query-request": 
+                return mock_objects_client
+            elif service_name == "response":
+                return flow_response
+            else:
+                return AsyncMock()
+        flow.side_effect = flow_router
         
         # Act
         await integration_processor.on_message(msg, consumer, flow)