Fixing tests

This commit is contained in:
Cyber MacGeddon 2025-09-30 13:44:18 +01:00
parent 02a908e46f
commit 48ac9960a7
2 changed files with 42 additions and 15 deletions

View file

@ -45,7 +45,7 @@ class TestEndToEndConfigurationFlow:
mock_message.triples = []
# Mock collection_exists to return True
with patch.object(processor.tg, 'collection_exists', return_value=True):
with patch('trustgraph.direct.cassandra_kg.KnowledgeGraph.collection_exists', return_value=True):
# This should create TrustGraph with environment config
await processor.store_triples(mock_message)
@ -149,7 +149,7 @@ class TestConfigurationPriorityEndToEnd:
mock_message.triples = []
# Mock collection_exists to return True
with patch.object(processor.tg, 'collection_exists', return_value=True):
with patch('trustgraph.direct.cassandra_kg.KnowledgeGraph.collection_exists', return_value=True):
await processor.store_triples(mock_message)
# Should use CLI parameters, not environment
@ -249,7 +249,7 @@ class TestNoBackwardCompatibilityEndToEnd:
mock_message.triples = []
# Mock collection_exists to return True
with patch.object(processor.tg, 'collection_exists', return_value=True):
with patch('trustgraph.direct.cassandra_kg.KnowledgeGraph.collection_exists', return_value=True):
await processor.store_triples(mock_message)
# Should use defaults since old parameters are not recognized
@ -307,7 +307,7 @@ class TestNoBackwardCompatibilityEndToEnd:
mock_message.triples = []
# Mock collection_exists to return True
with patch.object(processor.tg, 'collection_exists', return_value=True):
with patch('trustgraph.direct.cassandra_kg.KnowledgeGraph.collection_exists', return_value=True):
await processor.store_triples(mock_message)
# Should use new parameters, not old ones
@ -359,7 +359,7 @@ class TestMultipleHostsHandling:
mock_message.triples = []
# Mock collection_exists to return True
with patch.object(processor.tg, 'collection_exists', return_value=True):
with patch('trustgraph.direct.cassandra_kg.KnowledgeGraph.collection_exists', return_value=True):
await processor.store_triples(mock_message)
# Single host should be converted to list

View file

@ -23,14 +23,32 @@ class TestObjectsCassandraIntegration:
"""Mock Cassandra session for integration tests"""
session = MagicMock()
# Track if keyspaces have been created
created_keyspaces = set()
# Mock the execute method to return a valid result for keyspace checks
def execute_mock(query, *args, **kwargs):
result = MagicMock()
# For keyspace existence checks, return a row (keyspace exists)
if "system_schema.keyspaces" in str(query):
result.one.return_value = MagicMock() # Non-None means exists
query_str = str(query)
# Track keyspace creation
if "CREATE KEYSPACE" in query_str:
# Extract keyspace name from query
import re
match = re.search(r'CREATE KEYSPACE IF NOT EXISTS (\w+)', query_str)
if match:
created_keyspaces.add(match.group(1))
# For keyspace existence checks
if "system_schema.keyspaces" in query_str:
# Check if this keyspace was created
if args and args[0] in created_keyspaces:
result.one.return_value = MagicMock() # Exists
else:
result.one.return_value = None # Doesn't exist
else:
result.one.return_value = None
return result
session.execute = MagicMock(side_effect=execute_mock)
@ -96,7 +114,10 @@ class TestObjectsCassandraIntegration:
await processor.on_schema_config(config, version=1)
assert "customer_records" in processor.schemas
# Step 1.5: Create the collection first (simulate tg-set-collection)
await processor.create_collection("test_user", "import_2024")
# Step 2: Process an ExtractedObject
test_obj = ExtractedObject(
metadata=Metadata(
@ -115,10 +136,10 @@ class TestObjectsCassandraIntegration:
confidence=0.95,
source_span="Customer: John Doe..."
)
msg = MagicMock()
msg.value.return_value = test_obj
await processor.on_object(msg, None, None)
# Verify Cassandra interactions
@ -189,7 +210,11 @@ class TestObjectsCassandraIntegration:
await processor.on_schema_config(config, version=1)
assert len(processor.schemas) == 2
# Create collections first
await processor.create_collection("shop", "catalog")
await processor.create_collection("shop", "sales")
# Process objects for different schemas
product_obj = ExtractedObject(
metadata=Metadata(id="p1", user="shop", collection="catalog", metadata=[]),
@ -198,7 +223,7 @@ class TestObjectsCassandraIntegration:
confidence=0.9,
source_span="Product..."
)
order_obj = ExtractedObject(
metadata=Metadata(id="o1", user="shop", collection="sales", metadata=[]),
schema_name="orders",
@ -206,7 +231,7 @@ class TestObjectsCassandraIntegration:
confidence=0.85,
source_span="Order..."
)
# Process both objects
for obj in [product_obj, order_obj]:
msg = MagicMock()
@ -336,8 +361,10 @@ class TestObjectsCassandraIntegration:
)
# Make insert fail
mock_result = MagicMock()
mock_result.one.return_value = MagicMock() # Keyspace exists
mock_session.execute.side_effect = [
None, # keyspace creation succeeds
mock_result, # keyspace existence check succeeds
None, # table creation succeeds
Exception("Connection timeout") # insert fails
]