mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 12:41:02 +02:00
Add schema select + tests
This commit is contained in:
Cyber MacGeddon
parent
82334a5b69
commit
c5904b2180
4 changed files with 619 additions and 2 deletions
|
|
@ -0,0 +1,3 @@
|
|||
"""
|
||||
Unit and contract tests for structured-diag service
|
||||
"""
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue