From 48ac9960a78e535ce111f50365f9e07f9902821d Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 30 Sep 2025 13:44:18 +0100 Subject: [PATCH] Fixing tests --- .../test_cassandra_config_end_to_end.py | 10 ++-- .../test_objects_cassandra_integration.py | 47 +++++++++++++++---- 2 files changed, 42 insertions(+), 15 deletions(-) diff --git a/tests/integration/test_cassandra_config_end_to_end.py b/tests/integration/test_cassandra_config_end_to_end.py index 107282ab..e706b76a 100644 --- a/tests/integration/test_cassandra_config_end_to_end.py +++ b/tests/integration/test_cassandra_config_end_to_end.py @@ -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 diff --git a/tests/integration/test_objects_cassandra_integration.py b/tests/integration/test_objects_cassandra_integration.py index 6c5ca6f2..9a5d6056 100644 --- a/tests/integration/test_objects_cassandra_integration.py +++ b/tests/integration/test_objects_cassandra_integration.py @@ -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 ]