From 02a908e46f5eabdba4091727d5d4a3e7637a8a75 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 30 Sep 2025 13:37:28 +0100 Subject: [PATCH] Fixing tests --- .../test_cassandra_config_end_to_end.py | 38 ++++++++++++------- .../test_objects_cassandra_integration.py | 13 ++++++- .../storage/objects/cassandra/write.py | 7 +++- 3 files changed, 42 insertions(+), 16 deletions(-) diff --git a/tests/integration/test_cassandra_config_end_to_end.py b/tests/integration/test_cassandra_config_end_to_end.py index a14c521c..107282ab 100644 --- a/tests/integration/test_cassandra_config_end_to_end.py +++ b/tests/integration/test_cassandra_config_end_to_end.py @@ -29,23 +29,25 @@ class TestEndToEndConfigurationFlow: 'CASSANDRA_USERNAME': 'integration-user', 'CASSANDRA_PASSWORD': 'integration-pass' } - + mock_cluster_instance = MagicMock() mock_session = MagicMock() mock_cluster_instance.connect.return_value = mock_session mock_cluster.return_value = mock_cluster_instance - + with patch.dict(os.environ, env_vars, clear=True): processor = TriplesWriter(taskgroup=MagicMock()) - + # Create a mock message to trigger TrustGraph creation mock_message = MagicMock() mock_message.metadata.user = 'test_user' mock_message.metadata.collection = 'test_collection' mock_message.triples = [] - - # This should create TrustGraph with environment config - await processor.store_triples(mock_message) + + # Mock collection_exists to return True + with patch.object(processor.tg, 'collection_exists', return_value=True): + # This should create TrustGraph with environment config + await processor.store_triples(mock_message) # Verify Cluster was created with correct hosts mock_cluster.assert_called_once() @@ -145,8 +147,10 @@ class TestConfigurationPriorityEndToEnd: mock_message.metadata.user = 'test_user' mock_message.metadata.collection = 'test_collection' mock_message.triples = [] - - await processor.store_triples(mock_message) + + # Mock collection_exists to return True + with patch.object(processor.tg, 'collection_exists', return_value=True): + await processor.store_triples(mock_message) # Should use CLI parameters, not environment mock_cluster.assert_called_once() @@ -243,8 +247,10 @@ class TestNoBackwardCompatibilityEndToEnd: mock_message.metadata.user = 'legacy_user' mock_message.metadata.collection = 'legacy_collection' mock_message.triples = [] - - await processor.store_triples(mock_message) + + # Mock collection_exists to return True + with patch.object(processor.tg, 'collection_exists', return_value=True): + await processor.store_triples(mock_message) # Should use defaults since old parameters are not recognized mock_cluster.assert_called_once() @@ -299,8 +305,10 @@ class TestNoBackwardCompatibilityEndToEnd: mock_message.metadata.user = 'precedence_user' mock_message.metadata.collection = 'precedence_collection' mock_message.triples = [] - - await processor.store_triples(mock_message) + + # Mock collection_exists to return True + with patch.object(processor.tg, 'collection_exists', return_value=True): + await processor.store_triples(mock_message) # Should use new parameters, not old ones mock_cluster.assert_called_once() @@ -349,8 +357,10 @@ class TestMultipleHostsHandling: mock_message.metadata.user = 'single_user' mock_message.metadata.collection = 'single_collection' mock_message.triples = [] - - await processor.store_triples(mock_message) + + # Mock collection_exists to return True + with patch.object(processor.tg, 'collection_exists', return_value=True): + await processor.store_triples(mock_message) # Single host should be converted to list mock_cluster.assert_called_once() diff --git a/tests/integration/test_objects_cassandra_integration.py b/tests/integration/test_objects_cassandra_integration.py index 4ce86f74..6c5ca6f2 100644 --- a/tests/integration/test_objects_cassandra_integration.py +++ b/tests/integration/test_objects_cassandra_integration.py @@ -22,7 +22,18 @@ class TestObjectsCassandraIntegration: def mock_cassandra_session(self): """Mock Cassandra session for integration tests""" session = MagicMock() - session.execute = MagicMock() + + # 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 + else: + result.one.return_value = None + return result + + session.execute = MagicMock(side_effect=execute_mock) return session @pytest.fixture diff --git a/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py b/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py index fdd32a55..90578438 100644 --- a/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/objects/cassandra/write.py @@ -295,6 +295,8 @@ class Processor(FlowProcessor): try: self.session.execute(create_table_cql) + if keyspace not in self.known_tables: + self.known_tables[keyspace] = set() self.known_tables[keyspace].add(table_key) logger.info(f"Ensured table exists: {safe_keyspace}.{safe_table}") @@ -362,7 +364,8 @@ class Processor(FlowProcessor): WHERE keyspace_name = %s """ result = self.session.execute(check_keyspace_cql, (safe_keyspace,)) - if not result.one(): + # Check if result is None (mock case) or has no rows + if result is None or not result.one(): error_msg = ( f"Collection {obj.metadata.collection} does not exist. " f"Create it first with tg-set-collection." @@ -371,6 +374,8 @@ class Processor(FlowProcessor): raise ValueError(error_msg) # Cache it if it exists self.known_keyspaces.add(safe_keyspace) + if safe_keyspace not in self.known_tables: + self.known_tables[safe_keyspace] = set() # Get schema definition schema = self.schemas.get(obj.schema_name)