From c5904b2180251fb3989f6a41569e0876dfe16407 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 16 Sep 2025 23:20:44 +0100 Subject: [PATCH] Add schema select + tests --- .../test_structured_diag/__init__.py | 3 + .../test_schema_contracts.py | 258 +++++++++++++ .../test_schema_selection.py | 355 ++++++++++++++++++ .../retrieval/structured_diag/service.py | 5 +- 4 files changed, 619 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_retrieval/test_structured_diag/__init__.py create mode 100644 tests/unit/test_retrieval/test_structured_diag/test_schema_contracts.py create mode 100644 tests/unit/test_retrieval/test_structured_diag/test_schema_selection.py diff --git a/tests/unit/test_retrieval/test_structured_diag/__init__.py b/tests/unit/test_retrieval/test_structured_diag/__init__.py new file mode 100644 index 00000000..a900cbbb --- /dev/null +++ b/tests/unit/test_retrieval/test_structured_diag/__init__.py @@ -0,0 +1,3 @@ +""" +Unit and contract tests for structured-diag service +""" \ No newline at end of file diff --git a/tests/unit/test_retrieval/test_structured_diag/test_schema_contracts.py b/tests/unit/test_retrieval/test_structured_diag/test_schema_contracts.py new file mode 100644 index 00000000..99f66dc7 --- /dev/null +++ b/tests/unit/test_retrieval/test_structured_diag/test_schema_contracts.py @@ -0,0 +1,258 @@ +""" +Contract tests for structured-diag service schemas +""" + +import pytest +import json +from pulsar.schema import JsonSchema +from trustgraph.schema.services.diagnosis import ( + StructuredDataDiagnosisRequest, + StructuredDataDiagnosisResponse +) + + +class TestStructuredDiagnosisSchemaContract: + """Contract tests for structured diagnosis message schemas""" + + def test_request_schema_basic_fields(self): + """Test basic request schema fields""" + request = StructuredDataDiagnosisRequest( + operation="detect-type", + sample="test data" + ) + + assert request.operation == "detect-type" + assert request.sample == "test data" + assert request.type is None # Optional, defaults to None + assert request.schema_name is None # Optional, defaults to None + assert request.options is None # Optional, defaults to None + + def test_request_schema_all_operations(self): + """Test request schema supports all operations""" + operations = ["detect-type", "generate-descriptor", "diagnose", "schema-selection"] + + for op in operations: + request = StructuredDataDiagnosisRequest( + operation=op, + sample="test data" + ) + assert request.operation == op + + def test_request_schema_with_options(self): + """Test request schema with options""" + options = {"delimiter": ",", "has_header": "true"} + request = StructuredDataDiagnosisRequest( + operation="generate-descriptor", + sample="test data", + type="csv", + schema_name="products", + options=options + ) + + assert request.options == options + assert request.type == "csv" + assert request.schema_name == "products" + + def test_response_schema_basic_fields(self): + """Test basic response schema fields""" + response = StructuredDataDiagnosisResponse( + operation="detect-type", + detected_type="xml", + confidence=0.9, + error=None # Explicitly set to None + ) + + assert response.operation == "detect-type" + assert response.detected_type == "xml" + assert response.confidence == 0.9 + assert response.error is None + assert response.descriptor is None + assert response.metadata is None + assert response.schema_matches is None # New field, defaults to None + + def test_response_schema_with_error(self): + """Test response schema with error""" + from trustgraph.schema.core.primitives import Error + + error = Error( + type="ServiceError", + message="Service unavailable" + ) + response = StructuredDataDiagnosisResponse( + operation="schema-selection", + error=error + ) + + assert response.error == error + assert response.error.type == "ServiceError" + assert response.error.message == "Service unavailable" + + def test_response_schema_with_schema_matches(self): + """Test response schema with schema_matches array""" + matches = ["products", "inventory", "catalog"] + response = StructuredDataDiagnosisResponse( + operation="schema-selection", + schema_matches=matches + ) + + assert response.operation == "schema-selection" + assert response.schema_matches == matches + assert len(response.schema_matches) == 3 + + def test_response_schema_empty_schema_matches(self): + """Test response schema with empty schema_matches array""" + response = StructuredDataDiagnosisResponse( + operation="schema-selection", + schema_matches=[] + ) + + assert response.schema_matches == [] + assert isinstance(response.schema_matches, list) + + def test_response_schema_with_descriptor(self): + """Test response schema with descriptor""" + descriptor = { + "mapping": { + "field1": "column1", + "field2": "column2" + } + } + response = StructuredDataDiagnosisResponse( + operation="generate-descriptor", + descriptor=json.dumps(descriptor) + ) + + assert response.descriptor == json.dumps(descriptor) + parsed = json.loads(response.descriptor) + assert parsed["mapping"]["field1"] == "column1" + + def test_response_schema_with_metadata(self): + """Test response schema with metadata""" + metadata = { + "csv_options": json.dumps({"delimiter": ","}), + "field_count": "5" + } + response = StructuredDataDiagnosisResponse( + operation="diagnose", + metadata=metadata + ) + + assert response.metadata == metadata + assert response.metadata["field_count"] == "5" + + def test_schema_serialization(self): + """Test that schemas can be serialized and deserialized correctly""" + # Test request serialization + request = StructuredDataDiagnosisRequest( + operation="schema-selection", + sample="test data", + options={"key": "value"} + ) + + # Simulate Pulsar JsonSchema serialization + schema = JsonSchema(StructuredDataDiagnosisRequest) + serialized = schema.encode(request) + deserialized = schema.decode(serialized) + + assert deserialized.operation == request.operation + assert deserialized.sample == request.sample + assert deserialized.options == request.options + + def test_response_serialization_with_schema_matches(self): + """Test response serialization with schema_matches array""" + response = StructuredDataDiagnosisResponse( + operation="schema-selection", + schema_matches=["schema1", "schema2"], + confidence=0.85 + ) + + # Simulate Pulsar JsonSchema serialization + schema = JsonSchema(StructuredDataDiagnosisResponse) + serialized = schema.encode(response) + deserialized = schema.decode(serialized) + + assert deserialized.operation == response.operation + assert deserialized.schema_matches == response.schema_matches + assert deserialized.confidence == response.confidence + + def test_backwards_compatibility(self): + """Test that old clients can still use the service without schema_matches""" + # Old response without schema_matches should still work + response = StructuredDataDiagnosisResponse( + operation="detect-type", + detected_type="json", + confidence=0.95 + ) + + # Verify default value for new field + assert response.schema_matches is None # Defaults to None when not set + + # Verify old fields still work + assert response.detected_type == "json" + assert response.confidence == 0.95 + + def test_schema_selection_operation_contract(self): + """Test complete contract for schema-selection operation""" + # Request + request = StructuredDataDiagnosisRequest( + operation="schema-selection", + sample="product_id,name,price\n1,Widget,9.99" + ) + + assert request.operation == "schema-selection" + assert request.sample != "" + + # Response with matches + response = StructuredDataDiagnosisResponse( + operation="schema-selection", + schema_matches=["products", "inventory"] + ) + + assert response.operation == "schema-selection" + assert isinstance(response.schema_matches, list) + assert len(response.schema_matches) == 2 + assert all(isinstance(s, str) for s in response.schema_matches) + + # Response with error + from trustgraph.schema.core.primitives import Error + error_response = StructuredDataDiagnosisResponse( + operation="schema-selection", + error=Error(type="PromptServiceError", message="Service unavailable") + ) + + assert error_response.error is not None + assert error_response.schema_matches is None # Default None when not set + + def test_all_operations_supported(self): + """Verify all operations are properly supported in the contract""" + supported_operations = { + "detect-type": { + "required_request": ["sample"], + "expected_response": ["detected_type", "confidence"] + }, + "generate-descriptor": { + "required_request": ["sample", "type", "schema_name"], + "expected_response": ["descriptor"] + }, + "diagnose": { + "required_request": ["sample"], + "expected_response": ["detected_type", "confidence", "descriptor"] + }, + "schema-selection": { + "required_request": ["sample"], + "expected_response": ["schema_matches"] + } + } + + for operation, contract in supported_operations.items(): + # Test request creation + request_data = {"operation": operation} + for field in contract["required_request"]: + request_data[field] = "test_value" + + request = StructuredDataDiagnosisRequest(**request_data) + assert request.operation == operation + + # Test response creation + response = StructuredDataDiagnosisResponse(operation=operation) + assert response.operation == operation \ No newline at end of file diff --git a/tests/unit/test_retrieval/test_structured_diag/test_schema_selection.py b/tests/unit/test_retrieval/test_structured_diag/test_schema_selection.py new file mode 100644 index 00000000..2a18287d --- /dev/null +++ b/tests/unit/test_retrieval/test_structured_diag/test_schema_selection.py @@ -0,0 +1,355 @@ +""" +Unit tests for structured-diag service schema-selection operation +""" + +import pytest +import json +from unittest.mock import AsyncMock, MagicMock, patch +from trustgraph.retrieval.structured_diag.service import Processor +from trustgraph.schema.services.diagnosis import StructuredDataDiagnosisRequest, StructuredDataDiagnosisResponse +from trustgraph.schema import RowSchema, Field as SchemaField, Error + + +@pytest.fixture +def mock_schemas(): + """Create mock schemas for testing""" + schemas = { + "products": RowSchema( + name="products", + description="Product catalog schema", + fields=[ + SchemaField( + name="product_id", + type="string", + description="Product identifier", + required=True, + primary=True, + indexed=True + ), + SchemaField( + name="name", + type="string", + description="Product name", + required=True + ), + SchemaField( + name="price", + type="number", + description="Product price", + required=True + ) + ] + ), + "customers": RowSchema( + name="customers", + description="Customer database schema", + fields=[ + SchemaField( + name="customer_id", + type="string", + description="Customer identifier", + required=True, + primary=True + ), + SchemaField( + name="name", + type="string", + description="Customer name", + required=True + ), + SchemaField( + name="email", + type="string", + description="Customer email", + required=True + ) + ] + ), + "orders": RowSchema( + name="orders", + description="Order management schema", + fields=[ + SchemaField( + name="order_id", + type="string", + description="Order identifier", + required=True, + primary=True + ), + SchemaField( + name="customer_id", + type="string", + description="Customer identifier", + required=True + ), + SchemaField( + name="total", + type="number", + description="Order total", + required=True + ) + ] + ) + } + return schemas + + +@pytest.fixture +def service(mock_schemas): + """Create service instance with mock configuration""" + service = Processor( + taskgroup=MagicMock(), + id="test-processor" + ) + service.schemas = mock_schemas + return service + + +@pytest.fixture +def mock_flow(): + """Create mock flow with prompt service""" + flow = MagicMock() + prompt_request_flow = AsyncMock() + flow.return_value.request = prompt_request_flow + return flow, prompt_request_flow + + +@pytest.mark.asyncio +async def test_schema_selection_success(service, mock_flow): + """Test successful schema selection""" + flow, prompt_request_flow = mock_flow + + # Mock prompt service response with matching schemas + mock_response = MagicMock() + mock_response.error = None + mock_response.text = '["products", "orders"]' + prompt_request_flow.return_value = mock_response + + # Create request + request = StructuredDataDiagnosisRequest( + operation="schema-selection", + sample="product_id,name,price,quantity\nPROD001,Widget,19.99,5" + ) + + # Execute operation + response = await service.schema_selection_operation(request, flow) + + # Verify response + assert response.error is None + assert response.operation == "schema-selection" + assert response.schema_matches == ["products", "orders"] + + # Verify prompt service was called correctly + prompt_request_flow.assert_called_once() + call_args = prompt_request_flow.call_args[0][0] + assert call_args.id == "schema-selection" + + # Check that all schemas were passed to prompt + terms = call_args.terms + schemas_data = json.loads(terms["schemas"]) + assert len(schemas_data) == 3 # All 3 schemas + assert any(s["name"] == "products" for s in schemas_data) + assert any(s["name"] == "customers" for s in schemas_data) + assert any(s["name"] == "orders" for s in schemas_data) + + +@pytest.mark.asyncio +async def test_schema_selection_empty_response(service, mock_flow): + """Test handling of empty prompt service response""" + flow, prompt_request_flow = mock_flow + + # Mock empty response from prompt service + mock_response = MagicMock() + mock_response.error = None + mock_response.text = "" + prompt_request_flow.return_value = mock_response + + # Create request + request = StructuredDataDiagnosisRequest( + operation="schema-selection", + sample="test data" + ) + + # Execute operation + response = await service.schema_selection_operation(request, flow) + + # Verify error response + assert response.error is not None + assert response.error.type == "PromptServiceError" + assert "Empty response" in response.error.message + assert response.operation == "schema-selection" + + +@pytest.mark.asyncio +async def test_schema_selection_prompt_error(service, mock_flow): + """Test handling of prompt service error""" + flow, prompt_request_flow = mock_flow + + # Mock error response from prompt service + mock_response = MagicMock() + mock_response.error = Error( + type="ServiceError", + message="Prompt service unavailable" + ) + mock_response.text = None + prompt_request_flow.return_value = mock_response + + # Create request + request = StructuredDataDiagnosisRequest( + operation="schema-selection", + sample="test data" + ) + + # Execute operation + response = await service.schema_selection_operation(request, flow) + + # Verify error response + assert response.error is not None + assert response.error.type == "PromptServiceError" + assert "Failed to select schemas" in response.error.message + assert response.operation == "schema-selection" + + +@pytest.mark.asyncio +async def test_schema_selection_invalid_json(service, mock_flow): + """Test handling of invalid JSON response from prompt service""" + flow, prompt_request_flow = mock_flow + + # Mock invalid JSON response + mock_response = MagicMock() + mock_response.error = None + mock_response.text = "not valid json" + prompt_request_flow.return_value = mock_response + + # Create request + request = StructuredDataDiagnosisRequest( + operation="schema-selection", + sample="test data" + ) + + # Execute operation + response = await service.schema_selection_operation(request, flow) + + # Verify error response + assert response.error is not None + assert response.error.type == "ParseError" + assert "Failed to parse schema selection response" in response.error.message + assert response.operation == "schema-selection" + + +@pytest.mark.asyncio +async def test_schema_selection_non_array_response(service, mock_flow): + """Test handling of non-array JSON response from prompt service""" + flow, prompt_request_flow = mock_flow + + # Mock non-array JSON response + mock_response = MagicMock() + mock_response.error = None + mock_response.text = '{"schema": "products"}' # Object instead of array + prompt_request_flow.return_value = mock_response + + # Create request + request = StructuredDataDiagnosisRequest( + operation="schema-selection", + sample="test data" + ) + + # Execute operation + response = await service.schema_selection_operation(request, flow) + + # Verify error response + assert response.error is not None + assert response.error.type == "ParseError" + assert "Failed to parse schema selection response" in response.error.message + assert response.operation == "schema-selection" + + +@pytest.mark.asyncio +async def test_schema_selection_with_options(service, mock_flow): + """Test schema selection with additional options""" + flow, prompt_request_flow = mock_flow + + # Mock successful response + mock_response = MagicMock() + mock_response.error = None + mock_response.text = '["products"]' + prompt_request_flow.return_value = mock_response + + # Create request with options + request = StructuredDataDiagnosisRequest( + operation="schema-selection", + sample="test data", + options={"filter": "catalog", "confidence": "high"} + ) + + # Execute operation + response = await service.schema_selection_operation(request, flow) + + # Verify response + assert response.error is None + assert response.schema_matches == ["products"] + + # Verify options were passed to prompt + call_args = prompt_request_flow.call_args[0][0] + terms = call_args.terms + options = json.loads(terms["options"]) + assert options["filter"] == "catalog" + assert options["confidence"] == "high" + + +@pytest.mark.asyncio +async def test_schema_selection_exception_handling(service, mock_flow): + """Test handling of unexpected exceptions""" + flow, prompt_request_flow = mock_flow + + # Mock exception during prompt service call + prompt_request_flow.side_effect = Exception("Unexpected error") + + # Create request + request = StructuredDataDiagnosisRequest( + operation="schema-selection", + sample="test data" + ) + + # Execute operation + response = await service.schema_selection_operation(request, flow) + + # Verify error response + assert response.error is not None + assert response.error.type == "PromptServiceError" + assert "Failed to select schemas" in response.error.message + assert response.operation == "schema-selection" + + +@pytest.mark.asyncio +async def test_schema_selection_empty_schemas(service, mock_flow): + """Test schema selection with no schemas configured""" + flow, prompt_request_flow = mock_flow + + # Clear schemas + service.schemas = {} + + # Mock response (shouldn't be reached) + mock_response = MagicMock() + mock_response.error = None + mock_response.text = '[]' + prompt_request_flow.return_value = mock_response + + # Create request + request = StructuredDataDiagnosisRequest( + operation="schema-selection", + sample="test data" + ) + + # Execute operation + response = await service.schema_selection_operation(request, flow) + + # Should still succeed but with empty schemas array passed to prompt + assert response.error is None + assert response.schema_matches == [] + + # Verify empty schemas array was passed + call_args = prompt_request_flow.call_args[0][0] + terms = call_args.terms + schemas_data = json.loads(terms["schemas"]) + assert len(schemas_data) == 0 \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/retrieval/structured_diag/service.py b/trustgraph-flow/trustgraph/retrieval/structured_diag/service.py index 8d0bd647..fc45c71c 100644 --- a/trustgraph-flow/trustgraph/retrieval/structured_diag/service.py +++ b/trustgraph-flow/trustgraph/retrieval/structured_diag/service.py @@ -313,7 +313,7 @@ class Processor(FlowProcessor): """Handle schema-selection operation""" logger.info("Processing schema-selection operation") - # Prepare all schemas for the prompt + # Prepare all schemas for the prompt - match the original config format all_schemas = [] for schema_name, row_schema in self.schemas.items(): schema_info = { @@ -327,7 +327,8 @@ class Processor(FlowProcessor): "required": f.required, "primary_key": f.primary, "indexed": f.indexed, - "enum_values": f.enum_values if f.enum_values else [] + "enum": f.enum_values if f.enum_values else [], + "size": f.size if hasattr(f, 'size') else 0 } for f in row_schema.fields ]