From f89e002614cf0856841c6d8d1c52cc3c179a600e Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 9 Mar 2026 10:22:08 +0000 Subject: [PATCH] Fixing tests --- .../test_doc_embeddings_milvus_storage.py | 35 ++++++++++------- .../test_doc_embeddings_pinecone_storage.py | 22 +++++------ .../test_graph_embeddings_milvus_storage.py | 39 +++++++++++-------- .../test_row_embeddings_qdrant_storage.py | 8 ++-- 4 files changed, 57 insertions(+), 47 deletions(-) diff --git a/tests/unit/test_storage/test_doc_embeddings_milvus_storage.py b/tests/unit/test_storage/test_doc_embeddings_milvus_storage.py index 6530c020..4229e6a0 100644 --- a/tests/unit/test_storage/test_doc_embeddings_milvus_storage.py +++ b/tests/unit/test_storage/test_doc_embeddings_milvus_storage.py @@ -245,26 +245,31 @@ class TestMilvusDocEmbeddingsStorageProcessor: message.metadata = MagicMock() message.metadata.user = 'test_user' message.metadata.collection = 'test_collection' - - chunk = ChunkEmbeddings( - chunk_id="Document with mixed dimensions", - vectors=[ - [0.1, 0.2], # 2D vector - [0.3, 0.4, 0.5, 0.6], # 4D vector - [0.7, 0.8, 0.9] # 3D vector - ] + + # Each chunk has a single vector of different dimensions + chunk1 = ChunkEmbeddings( + chunk_id="chunk/doc/2d", + vector=[0.1, 0.2] # 2D vector ) - message.chunks = [chunk] - + chunk2 = ChunkEmbeddings( + chunk_id="chunk/doc/4d", + vector=[0.3, 0.4, 0.5, 0.6] # 4D vector + ) + chunk3 = ChunkEmbeddings( + chunk_id="chunk/doc/3d", + vector=[0.7, 0.8, 0.9] # 3D vector + ) + message.chunks = [chunk1, chunk2, chunk3] + await processor.store_document_embeddings(message) - + # Verify all vectors were inserted regardless of dimension with user/collection parameters expected_calls = [ - ([0.1, 0.2], "Document with mixed dimensions", 'test_user', 'test_collection'), - ([0.3, 0.4, 0.5, 0.6], "Document with mixed dimensions", 'test_user', 'test_collection'), - ([0.7, 0.8, 0.9], "Document with mixed dimensions", 'test_user', 'test_collection'), + ([0.1, 0.2], "chunk/doc/2d", 'test_user', 'test_collection'), + ([0.3, 0.4, 0.5, 0.6], "chunk/doc/4d", 'test_user', 'test_collection'), + ([0.7, 0.8, 0.9], "chunk/doc/3d", 'test_user', 'test_collection'), ] - + assert processor.vecstore.insert.call_count == 3 for i, (expected_vec, expected_doc, expected_user, expected_collection) in enumerate(expected_calls): actual_call = processor.vecstore.insert.call_args_list[i] diff --git a/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py b/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py index fc7c0a79..e5c1e91e 100644 --- a/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py +++ b/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py @@ -27,11 +27,11 @@ class TestPineconeDocEmbeddingsStorageProcessor: # Create test document embeddings chunk1 = ChunkEmbeddings( chunk=b"This is the first document chunk", - vectors=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + vector=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6] ) chunk2 = ChunkEmbeddings( chunk=b"This is the second document chunk", - vectors=[[0.7, 0.8, 0.9]] + vector=[0.7, 0.8, 0.9] ) message.chunks = [chunk1, chunk2] @@ -125,7 +125,7 @@ class TestPineconeDocEmbeddingsStorageProcessor: chunk = ChunkEmbeddings( chunk=b"Test document content", - vectors=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + vector=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6] ) message.chunks = [chunk] @@ -190,7 +190,7 @@ class TestPineconeDocEmbeddingsStorageProcessor: chunk = ChunkEmbeddings( chunk=b"Test document content", - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.chunks = [chunk] @@ -222,7 +222,7 @@ class TestPineconeDocEmbeddingsStorageProcessor: chunk = ChunkEmbeddings( chunk=b"", - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.chunks = [chunk] @@ -244,7 +244,7 @@ class TestPineconeDocEmbeddingsStorageProcessor: chunk = ChunkEmbeddings( chunk=None, - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.chunks = [chunk] @@ -266,7 +266,7 @@ class TestPineconeDocEmbeddingsStorageProcessor: chunk = ChunkEmbeddings( chunk=b"", # Empty bytes - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.chunks = [chunk] @@ -368,7 +368,7 @@ class TestPineconeDocEmbeddingsStorageProcessor: chunk = ChunkEmbeddings( chunk=b"Test document content", - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.chunks = [chunk] @@ -393,7 +393,7 @@ class TestPineconeDocEmbeddingsStorageProcessor: chunk = ChunkEmbeddings( chunk=b"Test document content", - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.chunks = [chunk] @@ -419,7 +419,7 @@ class TestPineconeDocEmbeddingsStorageProcessor: chunk = ChunkEmbeddings( chunk="Document with Unicode: éñ中文🚀".encode('utf-8'), - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.chunks = [chunk] @@ -447,7 +447,7 @@ class TestPineconeDocEmbeddingsStorageProcessor: large_content = "A" * 10000 # 10KB of content chunk = ChunkEmbeddings( chunk=large_content.encode('utf-8'), - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.chunks = [chunk] diff --git a/tests/unit/test_storage/test_graph_embeddings_milvus_storage.py b/tests/unit/test_storage/test_graph_embeddings_milvus_storage.py index 48b9c4f6..925cdca3 100644 --- a/tests/unit/test_storage/test_graph_embeddings_milvus_storage.py +++ b/tests/unit/test_storage/test_graph_embeddings_milvus_storage.py @@ -180,7 +180,7 @@ class TestMilvusGraphEmbeddingsStorageProcessor: ) empty_entity = EntityEmbeddings( entity=Term(type=LITERAL, value=''), - vectors=[[0.4, 0.5, 0.6]], + vector=[0.4, 0.5, 0.6], chunk_id='' ) none_entity = EntityEmbeddings( @@ -238,26 +238,31 @@ class TestMilvusGraphEmbeddingsStorageProcessor: message.metadata = MagicMock() message.metadata.user = 'test_user' message.metadata.collection = 'test_collection' - - entity = EntityEmbeddings( - entity=Term(type=IRI, iri='http://example.com/entity'), - vectors=[ - [0.1, 0.2], # 2D vector - [0.3, 0.4, 0.5, 0.6], # 4D vector - [0.7, 0.8, 0.9] # 3D vector - ] + + # Each entity has a single vector of different dimensions + entity1 = EntityEmbeddings( + entity=Term(type=IRI, iri='http://example.com/entity1'), + vector=[0.1, 0.2] # 2D vector ) - message.entities = [entity] - + entity2 = EntityEmbeddings( + entity=Term(type=IRI, iri='http://example.com/entity2'), + vector=[0.3, 0.4, 0.5, 0.6] # 4D vector + ) + entity3 = EntityEmbeddings( + entity=Term(type=IRI, iri='http://example.com/entity3'), + vector=[0.7, 0.8, 0.9] # 3D vector + ) + message.entities = [entity1, entity2, entity3] + await processor.store_graph_embeddings(message) - + # Verify all vectors were inserted regardless of dimension expected_calls = [ - ([0.1, 0.2], 'http://example.com/entity'), - ([0.3, 0.4, 0.5, 0.6], 'http://example.com/entity'), - ([0.7, 0.8, 0.9], 'http://example.com/entity'), + ([0.1, 0.2], 'http://example.com/entity1'), + ([0.3, 0.4, 0.5, 0.6], 'http://example.com/entity2'), + ([0.7, 0.8, 0.9], 'http://example.com/entity3'), ] - + assert processor.vecstore.insert.call_count == 3 for i, (expected_vec, expected_entity) in enumerate(expected_calls): actual_call = processor.vecstore.insert.call_args_list[i] @@ -278,7 +283,7 @@ class TestMilvusGraphEmbeddingsStorageProcessor: ) literal_entity = EntityEmbeddings( entity=Term(type=LITERAL, value='literal entity text'), - vectors=[[0.4, 0.5, 0.6]] + vector=[0.4, 0.5, 0.6] ) message.entities = [uri_entity, literal_entity] diff --git a/tests/unit/test_storage/test_row_embeddings_qdrant_storage.py b/tests/unit/test_storage/test_row_embeddings_qdrant_storage.py index b4c5a5b4..15697691 100644 --- a/tests/unit/test_storage/test_row_embeddings_qdrant_storage.py +++ b/tests/unit/test_storage/test_row_embeddings_qdrant_storage.py @@ -197,7 +197,7 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase): index_name='customer_id', index_value=['CUST001'], text='CUST001', - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) embeddings_msg = RowEmbeddings( @@ -255,7 +255,7 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase): index_name='name', index_value=['John Doe'], text='John Doe', - vectors=[[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]] + vector=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6] ) embeddings_msg = RowEmbeddings( @@ -299,7 +299,7 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase): index_name='id', index_value=['123'], text='123', - vectors=[] # Empty vectors + vector=[] # Empty vector ) embeddings_msg = RowEmbeddings( @@ -342,7 +342,7 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase): index_name='id', index_value=['123'], text='123', - vectors=[[0.1, 0.2]] + vector=[0.1, 0.2] ) embeddings_msg = RowEmbeddings(