apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-24 20:51:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Fixing tests

Browse source
This commit is contained in:
Cyber MacGeddon 2025-09-30 15:34:39 +01:00
parent d51b1ff7ba
commit a94ebe8377
7 changed files with 170 additions and 214 deletions
Download patch file Download diff file Expand all files Collapse all files

42
tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py
View file

@ -178,8 +178,8 @@ class TestPineconeDocEmbeddingsStorageProcessor:
assert calls[2][1]['vectors'][0]['metadata']['doc'] == "This is the second document chunk" assert calls[2][1]['vectors'][0]['metadata']['doc'] == "This is the second document chunk"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_document_embeddings_index_creation(self, processor): async def test_store_document_embeddings_index_validation(self, processor):
"""Test automatic index creation when index doesn't exist""" """Test that writing to non-existent index raises ValueError"""
message = MagicMock() message = MagicMock()
message.metadata = MagicMock() message.metadata = MagicMock()
message.metadata.user = 'test_user' message.metadata.user = 'test_user'
@ -191,25 +191,12 @@ class TestPineconeDocEmbeddingsStorageProcessor:
) )
message.chunks = [chunk] message.chunks = [chunk]
# Mock index doesn't exist initially # Mock index doesn't exist
processor.pinecone.has_index.return_value = False processor.pinecone.has_index.return_value = False
mock_index = MagicMock()
processor.pinecone.Index.return_value = mock_index
# Mock index creation with pytest.raises(ValueError, match="Collection .* does not exist"):
processor.pinecone.describe_index.return_value.status = {"ready": True}
with patch('uuid.uuid4', return_value='test-id'):
await processor.store_document_embeddings(message) await processor.store_document_embeddings(message)
# Verify index creation was called
expected_index_name = "d-test_user-test_collection"
processor.pinecone.create_index.assert_called_once()
create_call = processor.pinecone.create_index.call_args
assert create_call[1]['name'] == expected_index_name
assert create_call[1]['dimension'] == 3
assert create_call[1]['metric'] == "cosine"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_document_embeddings_empty_chunk(self, processor): async def test_store_document_embeddings_empty_chunk(self, processor):
"""Test storing document embeddings with empty chunk (should be skipped)""" """Test storing document embeddings with empty chunk (should be skipped)"""
@ -357,8 +344,8 @@ class TestPineconeDocEmbeddingsStorageProcessor:
mock_index.upsert.assert_not_called() mock_index.upsert.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_document_embeddings_index_creation_failure(self, processor): async def test_store_document_embeddings_validation_before_creation(self, processor):
"""Test handling of index creation failure""" """Test that validation error occurs before creation attempts"""
message = MagicMock() message = MagicMock()
message.metadata = MagicMock() message.metadata = MagicMock()
message.metadata.user = 'test_user' message.metadata.user = 'test_user'
@ -370,16 +357,15 @@ class TestPineconeDocEmbeddingsStorageProcessor:
) )
message.chunks = [chunk] message.chunks = [chunk]
# Mock index doesn't exist and creation fails # Mock index doesn't exist
processor.pinecone.has_index.return_value = False processor.pinecone.has_index.return_value = False
processor.pinecone.create_index.side_effect = Exception("Index creation failed")
with pytest.raises(Exception, match="Index creation failed"): with pytest.raises(ValueError, match="Collection .* does not exist"):
await processor.store_document_embeddings(message) await processor.store_document_embeddings(message)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_document_embeddings_index_creation_timeout(self, processor): async def test_store_document_embeddings_validates_before_timeout(self, processor):
"""Test handling of index creation timeout""" """Test that validation error occurs before timeout checks"""
message = MagicMock() message = MagicMock()
message.metadata = MagicMock() message.metadata = MagicMock()
message.metadata.user = 'test_user' message.metadata.user = 'test_user'
@ -391,13 +377,11 @@ class TestPineconeDocEmbeddingsStorageProcessor:
) )
message.chunks = [chunk] message.chunks = [chunk]
# Mock index doesn't exist and never becomes ready # Mock index doesn't exist
processor.pinecone.has_index.return_value = False processor.pinecone.has_index.return_value = False
processor.pinecone.describe_index.return_value.status = {"ready": False}
with patch('time.sleep'): # Speed up the test with pytest.raises(ValueError, match="Collection .* does not exist"):
with pytest.raises(RuntimeError, match="Gave up waiting for index creation"): await processor.store_document_embeddings(message)
await processor.store_document_embeddings(message)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_document_embeddings_unicode_content(self, processor): async def test_store_document_embeddings_unicode_content(self, processor):

48
tests/unit/test_storage/test_doc_embeddings_qdrant_storage.py
View file

@ -43,8 +43,6 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
# Verify processor attributes # Verify processor attributes
assert hasattr(processor, 'qdrant') assert hasattr(processor, 'qdrant')
assert processor.qdrant == mock_qdrant_instance assert processor.qdrant == mock_qdrant_instance
assert hasattr(processor, 'last_collection')
assert processor.last_collection is None
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__') @patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
@ -245,6 +243,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
# Arrange # Arrange
mock_base_init.return_value = None mock_base_init.return_value = None
mock_qdrant_instance = MagicMock() mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True # Collection exists
mock_qdrant_client.return_value = mock_qdrant_instance mock_qdrant_client.return_value = mock_qdrant_instance
config = { config = {
@ -273,12 +272,13 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
# Assert # Assert
# Should not call upsert for empty chunks # Should not call upsert for empty chunks
mock_qdrant_instance.upsert.assert_not_called() mock_qdrant_instance.upsert.assert_not_called()
mock_qdrant_instance.collection_exists.assert_not_called() # But collection_exists should be called for validation
mock_qdrant_instance.collection_exists.assert_called_once()
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__') @patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_collection_creation_when_not_exists(self, mock_base_init, mock_qdrant_client): async def test_collection_creation_when_not_exists(self, mock_base_init, mock_qdrant_client):
"""Test collection creation when it doesn't exist""" """Test that writing to non-existent collection raises ValueError"""
# Arrange # Arrange
mock_base_init.return_value = None mock_base_init.return_value = None
mock_qdrant_instance = MagicMock() mock_qdrant_instance = MagicMock()
@ -305,32 +305,18 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
mock_message.chunks = [mock_chunk] mock_message.chunks = [mock_chunk]
# Act # Act & Assert
await processor.store_document_embeddings(mock_message) with pytest.raises(ValueError, match="Collection .* does not exist"):
await processor.store_document_embeddings(mock_message)
# Assert
expected_collection = 'd_new_user_new_collection'
# Verify collection existence check and creation
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection)
mock_qdrant_instance.create_collection.assert_called_once()
# Verify create_collection was called with correct parameters
create_call_args = mock_qdrant_instance.create_collection.call_args
assert create_call_args[1]['collection_name'] == expected_collection
# Verify upsert was still called after collection creation
mock_qdrant_instance.upsert.assert_called_once()
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__') @patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_collection_creation_exception(self, mock_base_init, mock_qdrant_client): async def test_collection_creation_exception(self, mock_base_init, mock_qdrant_client):
"""Test collection creation handles exceptions""" """Test that validation error occurs before connection errors"""
# Arrange # Arrange
mock_base_init.return_value = None mock_base_init.return_value = None
mock_qdrant_instance = MagicMock() mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = False mock_qdrant_instance.collection_exists.return_value = False # Collection doesn't exist
mock_qdrant_instance.create_collection.side_effect = Exception("Qdrant connection failed")
mock_qdrant_client.return_value = mock_qdrant_instance mock_qdrant_client.return_value = mock_qdrant_instance
config = { config = {
@ -354,18 +340,21 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
mock_message.chunks = [mock_chunk] mock_message.chunks = [mock_chunk]
# Act & Assert # Act & Assert
with pytest.raises(Exception, match="Qdrant connection failed"): with pytest.raises(ValueError, match="Collection .* does not exist"):
await processor.store_document_embeddings(mock_message) await processor.store_document_embeddings(mock_message)
@patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__') @patch('trustgraph.base.DocumentEmbeddingsStoreService.__init__')
async def test_collection_caching_behavior(self, mock_base_init, mock_qdrant_client): @patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid')
"""Test collection caching with last_collection""" async def test_collection_validation_on_write(self, mock_uuid, mock_base_init, mock_qdrant_client):
"""Test collection validation checks collection exists before writing"""
# Arrange # Arrange
mock_base_init.return_value = None mock_base_init.return_value = None
mock_qdrant_instance = MagicMock() mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True mock_qdrant_instance.collection_exists.return_value = True
mock_qdrant_client.return_value = mock_qdrant_instance mock_qdrant_client.return_value = mock_qdrant_instance
mock_uuid.uuid4.return_value = MagicMock()
mock_uuid.uuid4.return_value.__str__ = MagicMock(return_value='test-uuid')
config = { config = {
'store_uri': 'http://localhost:6333', 'store_uri': 'http://localhost:6333',
@ -392,6 +381,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
# Reset mock to track second call # Reset mock to track second call
mock_qdrant_instance.reset_mock() mock_qdrant_instance.reset_mock()
mock_qdrant_instance.collection_exists.return_value = True
# Create second mock message with same dimensions # Create second mock message with same dimensions
mock_message2 = MagicMock() mock_message2 = MagicMock()
@ -409,11 +399,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase):
# Assert # Assert
expected_collection = 'd_cache_user_cache_collection' expected_collection = 'd_cache_user_cache_collection'
assert processor.last_collection == expected_collection
# Verify second call skipped existence check (cached) # Verify collection existence is checked on each write
mock_qdrant_instance.collection_exists.assert_not_called() mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection)
mock_qdrant_instance.create_collection.assert_not_called()
# But upsert should still be called # But upsert should still be called
mock_qdrant_instance.upsert.assert_called_once() mock_qdrant_instance.upsert.assert_called_once()

42
tests/unit/test_storage/test_graph_embeddings_pinecone_storage.py
View file

@ -178,8 +178,8 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
assert calls[2][1]['vectors'][0]['metadata']['entity'] == "entity2" assert calls[2][1]['vectors'][0]['metadata']['entity'] == "entity2"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_graph_embeddings_index_creation(self, processor): async def test_store_graph_embeddings_index_validation(self, processor):
"""Test automatic index creation when index doesn't exist""" """Test that writing to non-existent index raises ValueError"""
message = MagicMock() message = MagicMock()
message.metadata = MagicMock() message.metadata = MagicMock()
message.metadata.user = 'test_user' message.metadata.user = 'test_user'
@ -191,25 +191,12 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
) )
message.entities = [entity] message.entities = [entity]
# Mock index doesn't exist initially # Mock index doesn't exist
processor.pinecone.has_index.return_value = False processor.pinecone.has_index.return_value = False
mock_index = MagicMock()
processor.pinecone.Index.return_value = mock_index
# Mock index creation with pytest.raises(ValueError, match="Collection .* does not exist"):
processor.pinecone.describe_index.return_value.status = {"ready": True}
with patch('uuid.uuid4', return_value='test-id'):
await processor.store_graph_embeddings(message) await processor.store_graph_embeddings(message)
# Verify index creation was called
expected_index_name = "t-test_user-test_collection"
processor.pinecone.create_index.assert_called_once()
create_call = processor.pinecone.create_index.call_args
assert create_call[1]['name'] == expected_index_name
assert create_call[1]['dimension'] == 3
assert create_call[1]['metric'] == "cosine"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_graph_embeddings_empty_entity_value(self, processor): async def test_store_graph_embeddings_empty_entity_value(self, processor):
"""Test storing graph embeddings with empty entity value (should be skipped)""" """Test storing graph embeddings with empty entity value (should be skipped)"""
@ -328,8 +315,8 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
mock_index.upsert.assert_not_called() mock_index.upsert.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_graph_embeddings_index_creation_failure(self, processor): async def test_store_graph_embeddings_validation_before_creation(self, processor):
"""Test handling of index creation failure""" """Test that validation error occurs before any creation attempts"""
message = MagicMock() message = MagicMock()
message.metadata = MagicMock() message.metadata = MagicMock()
message.metadata.user = 'test_user' message.metadata.user = 'test_user'
@ -341,16 +328,15 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
) )
message.entities = [entity] message.entities = [entity]
# Mock index doesn't exist and creation fails # Mock index doesn't exist
processor.pinecone.has_index.return_value = False processor.pinecone.has_index.return_value = False
processor.pinecone.create_index.side_effect = Exception("Index creation failed")
with pytest.raises(Exception, match="Index creation failed"): with pytest.raises(ValueError, match="Collection .* does not exist"):
await processor.store_graph_embeddings(message) await processor.store_graph_embeddings(message)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_graph_embeddings_index_creation_timeout(self, processor): async def test_store_graph_embeddings_validates_before_timeout(self, processor):
"""Test handling of index creation timeout""" """Test that validation error occurs before timeout checks"""
message = MagicMock() message = MagicMock()
message.metadata = MagicMock() message.metadata = MagicMock()
message.metadata.user = 'test_user' message.metadata.user = 'test_user'
@ -362,13 +348,11 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
) )
message.entities = [entity] message.entities = [entity]
# Mock index doesn't exist and never becomes ready # Mock index doesn't exist
processor.pinecone.has_index.return_value = False processor.pinecone.has_index.return_value = False
processor.pinecone.describe_index.return_value.status = {"ready": False}
with patch('time.sleep'): # Speed up the test with pytest.raises(ValueError, match="Collection .* does not exist"):
with pytest.raises(RuntimeError, match="Gave up waiting for index creation"): await processor.store_graph_embeddings(message)
await processor.store_graph_embeddings(message)
def test_add_args_method(self): def test_add_args_method(self):
"""Test that add_args properly configures argument parser""" """Test that add_args properly configures argument parser"""

47
tests/unit/test_storage/test_graph_embeddings_qdrant_storage.py
View file

@ -43,13 +43,11 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
# Verify processor attributes # Verify processor attributes
assert hasattr(processor, 'qdrant') assert hasattr(processor, 'qdrant')
assert processor.qdrant == mock_qdrant_instance assert processor.qdrant == mock_qdrant_instance
assert hasattr(processor, 'last_collection')
assert processor.last_collection is None
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__') @patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
async def test_get_collection_creates_new_collection(self, mock_base_init, mock_qdrant_client): async def test_get_collection_validates_existence(self, mock_base_init, mock_qdrant_client):
"""Test get_collection creates a new collection when it doesn't exist""" """Test get_collection validates that collection exists"""
# Arrange # Arrange
mock_base_init.return_value = None mock_base_init.return_value = None
mock_qdrant_instance = MagicMock() mock_qdrant_instance = MagicMock()
@ -65,21 +63,9 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Act # Act & Assert
collection_name = processor.get_collection(dim=512, user='test_user', collection='test_collection') with pytest.raises(ValueError, match="Collection .* does not exist"):
processor.get_collection(user='test_user', collection='test_collection')
# Assert
expected_name = 't_test_user_test_collection'
assert collection_name == expected_name
assert processor.last_collection == expected_name
# Verify collection existence check and creation
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_name)
mock_qdrant_instance.create_collection.assert_called_once()
# Verify create_collection was called with correct parameters
create_call_args = mock_qdrant_instance.create_collection.call_args
assert create_call_args[1]['collection_name'] == expected_name
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid') @patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid')
@ -153,12 +139,11 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Act # Act
collection_name = processor.get_collection(dim=256, user='existing_user', collection='existing_collection') collection_name = processor.get_collection(user='existing_user', collection='existing_collection')
# Assert # Assert
expected_name = 't_existing_user_existing_collection' expected_name = 't_existing_user_existing_collection'
assert collection_name == expected_name assert collection_name == expected_name
assert processor.last_collection == expected_name
# Verify collection existence check was performed # Verify collection existence check was performed
mock_qdrant_instance.collection_exists.assert_called_once_with(expected_name) mock_qdrant_instance.collection_exists.assert_called_once_with(expected_name)
@ -167,8 +152,8 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__') @patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
async def test_get_collection_caches_last_collection(self, mock_base_init, mock_qdrant_client): async def test_get_collection_validates_on_each_call(self, mock_base_init, mock_qdrant_client):
"""Test get_collection skips checks when using same collection""" """Test get_collection validates collection existence on each call"""
# Arrange # Arrange
mock_base_init.return_value = None mock_base_init.return_value = None
mock_qdrant_instance = MagicMock() mock_qdrant_instance = MagicMock()
@ -185,32 +170,32 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# First call # First call
collection_name1 = processor.get_collection(dim=128, user='cache_user', collection='cache_collection') collection_name1 = processor.get_collection(user='cache_user', collection='cache_collection')
# Reset mock to track second call # Reset mock to track second call
mock_qdrant_instance.reset_mock() mock_qdrant_instance.reset_mock()
mock_qdrant_instance.collection_exists.return_value = True
# Act - Second call with same parameters # Act - Second call with same parameters
collection_name2 = processor.get_collection(dim=128, user='cache_user', collection='cache_collection') collection_name2 = processor.get_collection(user='cache_user', collection='cache_collection')
# Assert # Assert
expected_name = 't_cache_user_cache_collection' expected_name = 't_cache_user_cache_collection'
assert collection_name1 == expected_name assert collection_name1 == expected_name
assert collection_name2 == expected_name assert collection_name2 == expected_name
# Verify second call skipped existence check (cached) # Verify collection existence check happens on each call
mock_qdrant_instance.collection_exists.assert_not_called() mock_qdrant_instance.collection_exists.assert_called_once_with(expected_name)
mock_qdrant_instance.create_collection.assert_not_called() mock_qdrant_instance.create_collection.assert_not_called()
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.base.GraphEmbeddingsStoreService.__init__') @patch('trustgraph.base.GraphEmbeddingsStoreService.__init__')
async def test_get_collection_creation_exception(self, mock_base_init, mock_qdrant_client): async def test_get_collection_creation_exception(self, mock_base_init, mock_qdrant_client):
"""Test get_collection handles collection creation exceptions""" """Test get_collection raises ValueError when collection doesn't exist"""
# Arrange # Arrange
mock_base_init.return_value = None mock_base_init.return_value = None
mock_qdrant_instance = MagicMock() mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = False mock_qdrant_instance.collection_exists.return_value = False
mock_qdrant_instance.create_collection.side_effect = Exception("Qdrant connection failed")
mock_qdrant_client.return_value = mock_qdrant_instance mock_qdrant_client.return_value = mock_qdrant_instance
config = { config = {
@ -223,8 +208,8 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Act & Assert # Act & Assert
with pytest.raises(Exception, match="Qdrant connection failed"): with pytest.raises(ValueError, match="Collection .* does not exist"):
processor.get_collection(dim=512, user='error_user', collection='error_collection') processor.get_collection(user='error_user', collection='error_collection')
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid') @patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid')

12
tests/unit/test_storage/test_memgraph_user_collection_isolation.py
View file

@ -71,7 +71,9 @@ class TestMemgraphUserCollectionIsolation:
mock_message.metadata.user = "test_user" mock_message.metadata.user = "test_user"
mock_message.metadata.collection = "test_collection" mock_message.metadata.collection = "test_collection"
await processor.store_triples(mock_message) # Mock collection_exists to bypass validation in unit tests
with patch.object(processor, 'collection_exists', return_value=True):
await processor.store_triples(mock_message)
# Verify user/collection parameters were passed to all operations # Verify user/collection parameters were passed to all operations
# Should have: create_node (subject), create_node (object), relate_node = 3 calls # Should have: create_node (subject), create_node (object), relate_node = 3 calls
@ -117,7 +119,9 @@ class TestMemgraphUserCollectionIsolation:
mock_message.metadata.user = None mock_message.metadata.user = None
mock_message.metadata.collection = None mock_message.metadata.collection = None
await processor.store_triples(mock_message) # Mock collection_exists to bypass validation in unit tests
with patch.object(processor, 'collection_exists', return_value=True):
await processor.store_triples(mock_message)
# Verify defaults were used # Verify defaults were used
for call in mock_driver.execute_query.call_args_list: for call in mock_driver.execute_query.call_args_list:
@ -318,7 +322,9 @@ class TestMemgraphUserCollectionRegression:
message_user1.metadata.user = "user1" message_user1.metadata.user = "user1"
message_user1.metadata.collection = "collection1" message_user1.metadata.collection = "collection1"
await processor.store_triples(message_user1) # Mock collection_exists to bypass validation in unit tests
with patch.object(processor, 'collection_exists', return_value=True):
await processor.store_triples(message_user1)
# Verify that all storage operations included user1/collection1 parameters # Verify that all storage operations included user1/collection1 parameters
for call in mock_driver.execute_query.call_args_list: for call in mock_driver.execute_query.call_args_list:

22
tests/unit/test_storage/test_neo4j_user_collection_isolation.py
View file

@ -76,7 +76,9 @@ class TestNeo4jUserCollectionIsolation:
mock_summary.result_available_after = 10 mock_summary.result_available_after = 10
mock_driver.execute_query.return_value.summary = mock_summary mock_driver.execute_query.return_value.summary = mock_summary
await processor.store_triples(message) # Mock collection_exists to bypass validation in unit tests
with patch.object(processor, 'collection_exists', return_value=True):
await processor.store_triples(message)
# Verify nodes and relationships were created with user/collection properties # Verify nodes and relationships were created with user/collection properties
expected_calls = [ expected_calls = [
@ -142,7 +144,9 @@ class TestNeo4jUserCollectionIsolation:
mock_summary.result_available_after = 10 mock_summary.result_available_after = 10
mock_driver.execute_query.return_value.summary = mock_summary mock_driver.execute_query.return_value.summary = mock_summary
await processor.store_triples(message) # Mock collection_exists to bypass validation in unit tests
with patch.object(processor, 'collection_exists', return_value=True):
await processor.store_triples(message)
# Verify defaults were used # Verify defaults were used
mock_driver.execute_query.assert_any_call( mock_driver.execute_query.assert_any_call(
@ -274,9 +278,11 @@ class TestNeo4jUserCollectionIsolation:
mock_summary.result_available_after = 10 mock_summary.result_available_after = 10
mock_driver.execute_query.return_value.summary = mock_summary mock_driver.execute_query.return_value.summary = mock_summary
# Store data for both users # Mock collection_exists to bypass validation in unit tests
await processor.store_triples(message_user1) with patch.object(processor, 'collection_exists', return_value=True):
await processor.store_triples(message_user2) # Store data for both users
await processor.store_triples(message_user1)
await processor.store_triples(message_user2)
# Verify user1 data was stored with user1/coll1 # Verify user1 data was stored with user1/coll1
mock_driver.execute_query.assert_any_call( mock_driver.execute_query.assert_any_call(
@ -447,8 +453,10 @@ class TestNeo4jUserCollectionRegression:
mock_summary.result_available_after = 10 mock_summary.result_available_after = 10
mock_driver.execute_query.return_value.summary = mock_summary mock_driver.execute_query.return_value.summary = mock_summary
await processor.store_triples(message_user1) # Mock collection_exists to bypass validation in unit tests
await processor.store_triples(message_user2) with patch.object(processor, 'collection_exists', return_value=True):
await processor.store_triples(message_user1)
await processor.store_triples(message_user2)
# Verify two separate nodes were created with same URI but different user/collection # Verify two separate nodes were created with same URI but different user/collection
user1_node_call = call( user1_node_call = call(

7
tests/unit/test_storage/test_objects_cassandra_storage.py
View file

@ -293,15 +293,16 @@ class TestObjectsCassandraStorageLogic:
"""Test that secondary indexes are created for indexed fields""" """Test that secondary indexes are created for indexed fields"""
processor = MagicMock() processor = MagicMock()
processor.schemas = {} processor.schemas = {}
processor.known_keyspaces = set() processor.known_keyspaces = {"test_user"} # Pre-populate to skip validation query
processor.known_tables = {} processor.known_tables = {"test_user": set()} # Pre-populate
processor.session = MagicMock() processor.session = MagicMock()
processor.sanitize_name = Processor.sanitize_name.__get__(processor, Processor) processor.sanitize_name = Processor.sanitize_name.__get__(processor, Processor)
processor.sanitize_table = Processor.sanitize_table.__get__(processor, Processor) processor.sanitize_table = Processor.sanitize_table.__get__(processor, Processor)
processor.get_cassandra_type = Processor.get_cassandra_type.__get__(processor, Processor) processor.get_cassandra_type = Processor.get_cassandra_type.__get__(processor, Processor)
def mock_ensure_keyspace(keyspace): def mock_ensure_keyspace(keyspace):
processor.known_keyspaces.add(keyspace) processor.known_keyspaces.add(keyspace)
processor.known_tables[keyspace] = set() if keyspace not in processor.known_tables:
processor.known_tables[keyspace] = set()
processor.ensure_keyspace = mock_ensure_keyspace processor.ensure_keyspace = mock_ensure_keyspace
processor.ensure_table = Processor.ensure_table.__get__(processor, Processor) processor.ensure_table = Processor.ensure_table.__get__(processor, Processor)