From 9d95e5758b01f1f427dc7f31920fc41d99048328 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Thu, 4 Sep 2025 15:44:07 +0100
Subject: [PATCH] Fix tests to match recent changes

---
 tests/contract/test_message_contracts.py      |  2 +-
 .../test_structured_data_contracts.py         | 35 ++++++++-----------
 2 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/tests/contract/test_message_contracts.py b/tests/contract/test_message_contracts.py
index e4779d8b..972bf1f0 100644
--- a/tests/contract/test_message_contracts.py
+++ b/tests/contract/test_message_contracts.py
@@ -20,7 +20,7 @@ from trustgraph.schema import (
     GraphEmbeddings, EntityEmbeddings,
     Metadata, Field, RowSchema,
     StructuredDataSubmission, ExtractedObject,
-    NLPToStructuredQueryRequest, NLPToStructuredQueryResponse,
+    QuestionToStructuredQueryRequest, QuestionToStructuredQueryResponse,
     StructuredQueryRequest, StructuredQueryResponse,
     StructuredObjectEmbedding
 )
diff --git a/tests/contract/test_structured_data_contracts.py b/tests/contract/test_structured_data_contracts.py
index 43be9889..d3034ded 100644
--- a/tests/contract/test_structured_data_contracts.py
+++ b/tests/contract/test_structured_data_contracts.py
@@ -12,7 +12,7 @@ from typing import Dict, Any
 
 from trustgraph.schema import (
     StructuredDataSubmission, ExtractedObject,
-    NLPToStructuredQueryRequest, NLPToStructuredQueryResponse,
+    QuestionToStructuredQueryRequest, QuestionToStructuredQueryResponse,
     StructuredQueryRequest, StructuredQueryResponse,
     StructuredObjectEmbedding, Field, RowSchema,
     Metadata, Error, Value
@@ -146,23 +146,21 @@ class TestStructuredQueryServiceContracts:
     """Contract tests for structured query services"""
 
     def test_nlp_to_structured_query_request_contract(self):
-        """Test NLPToStructuredQueryRequest schema contract"""
+        """Test QuestionToStructuredQueryRequest schema contract"""
         # Act
-        request = NLPToStructuredQueryRequest(
-            natural_language_query="Show me all customers who registered last month",
-            max_results=100,
-            context_hints={"time_range": "last_month", "entity_type": "customer"}
+        request = QuestionToStructuredQueryRequest(
+            question="Show me all customers who registered last month",
+            max_results=100
         )
 
         # Assert
-        assert "customers" in request.natural_language_query
+        assert "customers" in request.question
         assert request.max_results == 100
-        assert request.context_hints["time_range"] == "last_month"
 
     def test_nlp_to_structured_query_response_contract(self):
-        """Test NLPToStructuredQueryResponse schema contract"""
+        """Test QuestionToStructuredQueryResponse schema contract"""
         # Act
-        response = NLPToStructuredQueryResponse(
+        response = QuestionToStructuredQueryResponse(
             error=None,
             graphql_query="query { customers(filter: {registered: {gte: \"2024-01-01\"}}) { id name email } }",
             variables={"start_date": "2024-01-01"},
@@ -180,15 +178,11 @@ class TestStructuredQueryServiceContracts:
         """Test StructuredQueryRequest schema contract"""
         # Act
         request = StructuredQueryRequest(
-            query="query GetCustomers($limit: Int) { customers(limit: $limit) { id name email } }",
-            variables={"limit": "10"},
-            operation_name="GetCustomers"
+            question="Show me customers with limit 10"
         )
 
         # Assert
-        assert "customers" in request.query
-        assert request.variables["limit"] == "10"
-        assert request.operation_name == "GetCustomers"
+        assert "customers" in request.question
 
     def test_structured_query_response_contract(self):
         """Test StructuredQueryResponse schema contract"""
@@ -291,11 +285,10 @@ class TestStructuredDataSerializationContracts:
         """Test NLP query request/response serialization contract"""
         # Test request
         request_data = {
-            "natural_language_query": "test query",
-            "max_results": 10,
-            "context_hints": {}
+            "question": "test query",
+            "max_results": 10
         }
-        assert serialize_deserialize_test(NLPToStructuredQueryRequest, request_data)
+        assert serialize_deserialize_test(QuestionToStructuredQueryRequest, request_data)
 
         # Test response
         response_data = {
@@ -305,4 +298,4 @@ class TestStructuredDataSerializationContracts:
             "detected_schemas": ["test"],
             "confidence": 0.9
         }
-        assert serialize_deserialize_test(NLPToStructuredQueryResponse, response_data)
\ No newline at end of file
+        assert serialize_deserialize_test(QuestionToStructuredQueryResponse, response_data)
\ No newline at end of file