From 1dc5ce2cb6dc1123d3a7014b8a73ce464e8fd3d1 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 5 Sep 2025 08:22:51 +0100 Subject: [PATCH] Fix tests --- .../integration/test_nlp_query_integration.py | 39 +++++-- .../test_structured_query_integration.py | 106 ++++++++++++++---- 2 files changed, 115 insertions(+), 30 deletions(-) diff --git a/tests/integration/test_nlp_query_integration.py b/tests/integration/test_nlp_query_integration.py index 83e56ac4..a7ee9d87 100644 --- a/tests/integration/test_nlp_query_integration.py +++ b/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 diff --git a/tests/integration/test_structured_query_integration.py b/tests/integration/test_structured_query_integration.py index 72ae1e77..b06aebc4 100644 --- a/tests/integration/test_structured_query_integration.py +++ b/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)