Fixing tests

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
+            # Mock collection_exists to return True
-            await processor.store_triples(mock_message)
+            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()

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

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)