apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-24 12:41:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Fix tests

Browse source
This commit is contained in:
Cyber MacGeddon 2025-09-05 08:26:42 +01:00
parent 1dc5ce2cb6
commit 86f0d41553
2 changed files with 115 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

44
tests/unit/test_retrieval/test_nlp_query.py
View file

@ -87,14 +87,18 @@ class TestNLPQueryProcessor:
error=None error=None
) )
processor.client.return_value.request = AsyncMock(return_value=mock_response) # Mock flow context
flow = MagicMock()
mock_prompt_service = AsyncMock()
mock_prompt_service.request = AsyncMock(return_value=mock_response)
flow.side_effect = lambda service_name: mock_prompt_service if service_name == "prompt-request" else AsyncMock()
# Act # Act
result = await processor.phase1_select_schemas(question) result = await processor.phase1_select_schemas(question, flow)
# Assert # Assert
assert result == expected_schemas assert result == expected_schemas
processor.client.assert_called_once_with("prompt-request") mock_prompt_service.request.assert_called_once()
async def test_phase1_select_schemas_prompt_error(self, processor): async def test_phase1_select_schemas_prompt_error(self, processor):
"""Test schema selection with prompt service error""" """Test schema selection with prompt service error"""
@ -103,11 +107,15 @@ class TestNLPQueryProcessor:
error = Error(type="prompt-error", message="Template not found") error = Error(type="prompt-error", message="Template not found")
mock_response = PromptResponse(text="", error=error) mock_response = PromptResponse(text="", error=error)
processor.client.return_value.request = AsyncMock(return_value=mock_response) # Mock flow context
flow = MagicMock()
mock_prompt_service = AsyncMock()
mock_prompt_service.request = AsyncMock(return_value=mock_response)
flow.side_effect = lambda service_name: mock_prompt_service if service_name == "prompt-request" else AsyncMock()
# Act & Assert # Act & Assert
with pytest.raises(Exception, match="Prompt service error"): with pytest.raises(Exception, match="Prompt service error"):
await processor.phase1_select_schemas(question) await processor.phase1_select_schemas(question, flow)
async def test_phase2_generate_graphql_success(self, processor): async def test_phase2_generate_graphql_success(self, processor):
"""Test successful GraphQL generation (Phase 2)""" """Test successful GraphQL generation (Phase 2)"""
@ -125,14 +133,18 @@ class TestNLPQueryProcessor:
error=None error=None
) )
processor.client.return_value.request = AsyncMock(return_value=mock_response) # Mock flow context
flow = MagicMock()
mock_prompt_service = AsyncMock()
mock_prompt_service.request = AsyncMock(return_value=mock_response)
flow.side_effect = lambda service_name: mock_prompt_service if service_name == "prompt-request" else AsyncMock()
# Act # Act
result = await processor.phase2_generate_graphql(question, selected_schemas) result = await processor.phase2_generate_graphql(question, selected_schemas, flow)
# Assert # Assert
assert result == expected_result assert result == expected_result
processor.client.assert_called_once_with("prompt-request") mock_prompt_service.request.assert_called_once()
async def test_phase2_generate_graphql_prompt_error(self, processor): async def test_phase2_generate_graphql_prompt_error(self, processor):
"""Test GraphQL generation with prompt service error""" """Test GraphQL generation with prompt service error"""
@ -142,11 +154,15 @@ class TestNLPQueryProcessor:
error = Error(type="prompt-error", message="Generation failed") error = Error(type="prompt-error", message="Generation failed")
mock_response = PromptResponse(text="", error=error) mock_response = PromptResponse(text="", error=error)
processor.client.return_value.request = AsyncMock(return_value=mock_response) # Mock flow context
flow = MagicMock()
mock_prompt_service = AsyncMock()
mock_prompt_service.request = AsyncMock(return_value=mock_response)
flow.side_effect = lambda service_name: mock_prompt_service if service_name == "prompt-request" else AsyncMock()
# Act & Assert # Act & Assert
with pytest.raises(Exception, match="Prompt service error"): with pytest.raises(Exception, match="Prompt service error"):
await processor.phase2_generate_graphql(question, selected_schemas) await processor.phase2_generate_graphql(question, selected_schemas, flow)
async def test_on_message_full_flow_success(self, processor): async def test_on_message_full_flow_success(self, processor):
"""Test complete message processing flow""" """Test complete message processing flow"""
@ -181,16 +197,18 @@ class TestNLPQueryProcessor:
error=None error=None
) )
# Set up mock to return different responses for each call # Mock flow context to return prompt service responses
processor.client.return_value.request = AsyncMock( mock_prompt_service = AsyncMock()
mock_prompt_service.request = AsyncMock(
side_effect=[phase1_response, phase2_response] side_effect=[phase1_response, phase2_response]
) )
flow.side_effect = lambda service_name: mock_prompt_service if service_name == "prompt-request" else flow_response if service_name == "response" else AsyncMock()
# Act # Act
await processor.on_message(msg, consumer, flow) await processor.on_message(msg, consumer, flow)
# Assert # Assert
assert processor.client.return_value.request.call_count == 2 assert mock_prompt_service.request.call_count == 2
flow_response.send.assert_called_once() flow_response.send.assert_called_once()
# Verify response structure # Verify response structure

104
tests/unit/test_retrieval/test_structured_query.py
View file

@ -80,9 +80,17 @@ class TestStructuredQueryProcessor:
mock_objects_client = AsyncMock() mock_objects_client = AsyncMock()
mock_objects_client.request.return_value = objects_response mock_objects_client.request.return_value = objects_response
processor.client.side_effect = lambda name: ( # Mock flow context to route to appropriate services
mock_nlp_client if name == "nlp-query-request" else mock_objects_client 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 # Act
await processor.on_message(msg, consumer, flow) await processor.on_message(msg, consumer, flow)
@ -101,7 +109,7 @@ class TestStructuredQueryProcessor:
assert isinstance(objects_call_args, ObjectsQueryRequest) assert isinstance(objects_call_args, ObjectsQueryRequest)
assert objects_call_args.query == 'query { customers(where: {state: {eq: "NY"}}) { id name email } }' assert objects_call_args.query == 'query { customers(where: {state: {eq: "NY"}}) { id name email } }'
assert objects_call_args.variables == {"state": "NY"} assert objects_call_args.variables == {"state": "NY"}
assert objects_call_args.user == "default" assert objects_call_args.user == "trustgraph"
assert objects_call_args.collection == "default" assert objects_call_args.collection == "default"
# Verify response # Verify response
@ -142,7 +150,15 @@ class TestStructuredQueryProcessor:
mock_nlp_client = AsyncMock() mock_nlp_client = AsyncMock()
mock_nlp_client.request.return_value = nlp_response mock_nlp_client.request.return_value = nlp_response
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 # Act
await processor.on_message(msg, consumer, flow) await processor.on_message(msg, consumer, flow)
@ -185,7 +201,15 @@ class TestStructuredQueryProcessor:
mock_nlp_client = AsyncMock() mock_nlp_client = AsyncMock()
mock_nlp_client.request.return_value = nlp_response mock_nlp_client.request.return_value = nlp_response
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 # Act
await processor.on_message(msg, consumer, flow) await processor.on_message(msg, consumer, flow)
@ -237,9 +261,17 @@ class TestStructuredQueryProcessor:
mock_objects_client = AsyncMock() mock_objects_client = AsyncMock()
mock_objects_client.request.return_value = objects_response mock_objects_client.request.return_value = objects_response
processor.client.side_effect = lambda name: ( # Mock flow context to route to appropriate services
mock_nlp_client if name == "nlp-query-request" else mock_objects_client 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 # Act
await processor.on_message(msg, consumer, flow) await processor.on_message(msg, consumer, flow)
@ -300,9 +332,17 @@ class TestStructuredQueryProcessor:
mock_objects_client = AsyncMock() mock_objects_client = AsyncMock()
mock_objects_client.request.return_value = objects_response mock_objects_client.request.return_value = objects_response
processor.client.side_effect = lambda name: ( # Mock flow context to route to appropriate services
mock_nlp_client if name == "nlp-query-request" else mock_objects_client 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 # Act
await processor.on_message(msg, consumer, flow) await processor.on_message(msg, consumer, flow)
@ -370,9 +410,17 @@ class TestStructuredQueryProcessor:
mock_objects_client = AsyncMock() mock_objects_client = AsyncMock()
mock_objects_client.request.return_value = objects_response mock_objects_client.request.return_value = objects_response
processor.client.side_effect = lambda name: ( # Mock flow context to route to appropriate services
mock_nlp_client if name == "nlp-query-request" else mock_objects_client 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 # Act
await processor.on_message(msg, consumer, flow) await processor.on_message(msg, consumer, flow)
@ -427,9 +475,17 @@ class TestStructuredQueryProcessor:
mock_objects_client = AsyncMock() mock_objects_client = AsyncMock()
mock_objects_client.request.return_value = objects_response mock_objects_client.request.return_value = objects_response
processor.client.side_effect = lambda name: ( # Mock flow context to route to appropriate services
mock_nlp_client if name == "nlp-query-request" else mock_objects_client 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 # Act
await processor.on_message(msg, consumer, flow) await processor.on_message(msg, consumer, flow)
@ -457,10 +513,18 @@ class TestStructuredQueryProcessor:
flow_response = AsyncMock() flow_response = AsyncMock()
flow.return_value = flow_response flow.return_value = flow_response
# Mock client to raise exception # Mock flow context to raise exception
mock_client = AsyncMock() mock_client = AsyncMock()
mock_client.request.side_effect = Exception("Network timeout") mock_client.request.side_effect = Exception("Network timeout")
processor.client.return_value = mock_client
def flow_router(service_name):
if service_name == "nlp-query-request":
return mock_client
elif service_name == "response":
return flow_response
else:
return AsyncMock()
flow.side_effect = flow_router
# Act # Act
await processor.on_message(msg, consumer, flow) await processor.on_message(msg, consumer, flow)