Fixing tests

This commit is contained in:
Cyber MacGeddon 2025-09-08 17:52:22 +01:00
parent d03eb194d4
commit b32aa929dc
4 changed files with 21 additions and 13 deletions

View file

@ -120,7 +120,8 @@ Args: {
# Verify structured query was called
mock_structured_client.structured_query.assert_called_once()
call_args = mock_structured_client.structured_query.call_args
question_arg = call_args[0][0] # positional argument
# Check keyword arguments
question_arg = call_args.kwargs.get("question") or call_args[1].get("question")
assert "customers" in question_arg.lower()
assert "new york" in question_arg.lower()
@ -202,8 +203,9 @@ Args: {
# Agent should handle the error gracefully
assert any(isinstance(resp, AgentResponse) for resp in responses)
# The tool should have returned an error response that contains error info
structured_query_call_args = mock_structured_client.structured_query.call_args[0]
assert "table" in structured_query_call_args[0].lower() or "exist" in structured_query_call_args[0].lower()
call_args = mock_structured_client.structured_query.call_args
question_arg = call_args.kwargs.get("question") or call_args[1].get("question")
assert "table" in question_arg.lower() or "exist" in question_arg.lower()
@pytest.mark.asyncio
async def test_agent_multi_step_structured_query_reasoning(self, agent_processor, structured_query_tool_config):
@ -278,8 +280,9 @@ Args: {
assert any(isinstance(resp, AgentResponse) for resp in responses)
# Verify the structured query was called with customer-related question
call_args = mock_structured_client.structured_query.call_args[0]
assert "california" in call_args[0].lower()
call_args = mock_structured_client.structured_query.call_args
question_arg = call_args.kwargs.get("question") or call_args[1].get("question")
assert "california" in question_arg.lower()
@pytest.mark.asyncio
async def test_agent_structured_query_with_collection_parameter(self, agent_processor):
@ -367,8 +370,9 @@ Args: {
assert any(isinstance(resp, AgentResponse) for resp in responses)
# Check the query was about sales/transactions
call_args = mock_structured_client.structured_query.call_args[0]
assert "sales" in call_args[0].lower() or "transactions" in call_args[0].lower()
call_args = mock_structured_client.structured_query.call_args
question_arg = call_args.kwargs.get("question") or call_args[1].get("question")
assert "sales" in question_arg.lower() or "transactions" in question_arg.lower()
@pytest.mark.asyncio
async def test_agent_structured_query_tool_argument_validation(self, agent_processor, structured_query_tool_config):
@ -473,5 +477,6 @@ Args: {
assert any(isinstance(resp, AgentResponse) for resp in responses)
# Check that the query was about customer information
call_args = mock_structured_client.structured_query.call_args[0]
assert "customer" in call_args[0].lower()
call_args = mock_structured_client.structured_query.call_args
question_arg = call_args.kwargs.get("question") or call_args[1].get("question")
assert "customer" in question_arg.lower()

View file

@ -41,7 +41,9 @@ class TestStructuredQueryServiceIntegration:
"""Test complete structured query processing pipeline"""
# Arrange - Create realistic query request
request = StructuredQueryRequest(
question="Show me all customers from California who have made purchases over $500"
question="Show me all customers from California who have made purchases over $500",
user="trustgraph",
collection="default"
)
msg = MagicMock()

View file

@ -148,7 +148,8 @@ class Processor(AgentService):
elif impl_id == "structured-query":
impl = functools.partial(
StructuredQueryImpl,
collection=data.get("collection")
collection=data.get("collection"),
user=None # User will be provided dynamically via context
)
arguments = StructuredQueryImpl.get_arguments()
else:

View file

@ -29,9 +29,9 @@ def filter_tools_by_group_and_state(
"""
# Apply defaults as specified in tech spec
if requested_groups is None:
if requested_groups is None or requested_groups == []:
requested_groups = ["default"]
if current_state is None:
if current_state is None or current_state == "":
current_state = "undefined"
logger.info(f"Filtering tools with groups={requested_groups}, state={current_state}")