apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-05-03 12:22:37 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fixing tests

Browse source
This commit is contained in:
Cyber MacGeddon 2026-03-09 10:48:05 +00:00
parent 2670c5f80e
commit d211dcb58d
2 changed files with 43 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -70,7 +70,7 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
mock_entity = MagicMock()
mock_entity.entity.type = IRI
mock_entity.entity.iri = 'test_entity'
mock_entity.vectors = [[0.1, 0.2, 0.3]] # Single vector with 3 dimensions
mock_entity.vector = [0.1, 0.2, 0.3] # Single vector with 3 dimensions
mock_message.entities = [mock_entity]
@ -124,12 +124,12 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
mock_entity1 = MagicMock()
mock_entity1.entity.type = IRI
mock_entity1.entity.iri = 'entity_one'
mock_entity1.vectors = [[0.1, 0.2]]
mock_entity1.vector = [0.1, 0.2]
mock_entity2 = MagicMock()
mock_entity2.entity.type = IRI
mock_entity2.entity.iri = 'entity_two'
mock_entity2.vectors = [[0.3, 0.4]]
mock_entity2.vector = [0.3, 0.4]
mock_message.entities = [mock_entity1, mock_entity2]
@ -157,14 +157,14 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.graph_embeddings.qdrant.write.uuid')
async def test_store_graph_embeddings_multiple_vectors_per_entity(self, mock_uuid, mock_qdrant_client):
"""Test storing graph embeddings with multiple vectors per entity"""
async def test_store_graph_embeddings_three_entities(self, mock_uuid, mock_qdrant_client):
"""Test storing graph embeddings with three entities"""
# Arrange
mock_qdrant_instance = MagicMock()
mock_qdrant_instance.collection_exists.return_value = True
mock_qdrant_client.return_value = mock_qdrant_instance
mock_uuid.uuid4.return_value.return_value = 'test-uuid'
config = {
'store_uri': 'http://localhost:6333',
'api_key': 'test-api-key',
@ -177,42 +177,48 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
# Add collection to known_collections (simulates config push)
processor.known_collections[('vector_user', 'vector_collection')] = {}
# Create mock message with entity having multiple vectors
# Create mock message with three entities
mock_message = MagicMock()
mock_message.metadata.user = 'vector_user'
mock_message.metadata.collection = 'vector_collection'
mock_entity = MagicMock()
mock_entity.entity.type = IRI
mock_entity.entity.iri = 'multi_vector_entity'
mock_entity.vectors = [
[0.1, 0.2, 0.3],
[0.4, 0.5, 0.6],
[0.7, 0.8, 0.9]
]
mock_message.entities = [mock_entity]
mock_entity1 = MagicMock()
mock_entity1.entity.type = IRI
mock_entity1.entity.iri = 'entity_one'
mock_entity1.vector = [0.1, 0.2, 0.3]
mock_entity2 = MagicMock()
mock_entity2.entity.type = IRI
mock_entity2.entity.iri = 'entity_two'
mock_entity2.vector = [0.4, 0.5, 0.6]
mock_entity3 = MagicMock()
mock_entity3.entity.type = IRI
mock_entity3.entity.iri = 'entity_three'
mock_entity3.vector = [0.7, 0.8, 0.9]
mock_message.entities = [mock_entity1, mock_entity2, mock_entity3]
# Act
await processor.store_graph_embeddings(mock_message)
# Assert
# Should be called 3 times (once per vector)
# Should be called 3 times (once per entity)
assert mock_qdrant_instance.upsert.call_count == 3
# Verify all vectors were processed
# Verify all entities were processed
upsert_calls = mock_qdrant_instance.upsert.call_args_list
expected_vectors = [
[0.1, 0.2, 0.3],
[0.4, 0.5, 0.6],
[0.7, 0.8, 0.9]
expected_data = [
([0.1, 0.2, 0.3], 'entity_one'),
([0.4, 0.5, 0.6], 'entity_two'),
([0.7, 0.8, 0.9], 'entity_three')
]
for i, call in enumerate(upsert_calls):
point = call[1]['points'][0]
assert point.vector == expected_vectors[i]
assert point.payload['entity'] == 'multi_vector_entity'
assert point.vector == expected_data[i][0]
assert point.payload['entity'] == expected_data[i][1]
@patch('trustgraph.storage.graph_embeddings.qdrant.write.QdrantClient')
async def test_store_graph_embeddings_empty_entity_value(self, mock_qdrant_client):
@ -238,11 +244,11 @@ class TestQdrantGraphEmbeddingsStorage(IsolatedAsyncioTestCase):
mock_entity_empty = MagicMock()
mock_entity_empty.entity.type = LITERAL
mock_entity_empty.entity.value = "" # Empty string
mock_entity_empty.vectors = [[0.1, 0.2]]
mock_entity_empty.vector = [0.1, 0.2]
mock_entity_none = MagicMock()
mock_entity_none.entity = None # None entity
mock_entity_none.vectors = [[0.3, 0.4]]
mock_entity_none.vector = [0.3, 0.4]
mock_message.entities = [mock_entity_empty, mock_entity_none]

10
tests/unit/test_storage/test_row_embeddings_qdrant_storage.py
View file

@ -227,8 +227,8 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase):
@patch('trustgraph.storage.row_embeddings.qdrant.write.QdrantClient')
@patch('trustgraph.storage.row_embeddings.qdrant.write.uuid')
async def test_on_embeddings_multiple_vectors(self, mock_uuid, mock_qdrant_client):
"""Test processing embeddings with multiple vectors"""
async def test_on_embeddings_single_vector(self, mock_uuid, mock_qdrant_client):
"""Test processing embeddings with a single vector"""
from trustgraph.storage.row_embeddings.qdrant.write import Processor
from trustgraph.schema import RowEmbeddings, RowIndexEmbedding
@ -250,7 +250,7 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase):
metadata.collection = 'test_collection'
metadata.id = 'doc-123'
# Embedding with multiple vectors
# Embedding with a single 6D vector
embedding = RowIndexEmbedding(
index_name='name',
index_value=['John Doe'],
@ -269,8 +269,8 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase):
await processor.on_embeddings(mock_msg, MagicMock(), MagicMock())
# Should be called 3 times (once per vector)
assert mock_qdrant_instance.upsert.call_count == 3
# Should be called once for the single embedding
assert mock_qdrant_instance.upsert.call_count == 1
@patch('trustgraph.storage.row_embeddings.qdrant.write.QdrantClient')
async def test_on_embeddings_skips_empty_vectors(self, mock_qdrant_client):