mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-23 20:21:03 +02:00
Fixing tests
This commit is contained in:
parent
02a908e46f
commit
48ac9960a7
2 changed files with 42 additions and 15 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue