Fixing tests

This commit is contained in:
Cyber MacGeddon 2026-03-09 10:24:53 +00:00
parent f89e002614
commit 8d25f4d4a8
3 changed files with 39 additions and 33 deletions

View file

@ -98,7 +98,7 @@ def sample_graph_embeddings():
entities=[ entities=[
EntityEmbeddings( EntityEmbeddings(
entity=Term(type=IRI, iri="http://example.org/john"), entity=Term(type=IRI, iri="http://example.org/john"),
vectors=[[0.1, 0.2, 0.3]] vector=[0.1, 0.2, 0.3]
) )
] ]
) )

View file

@ -286,37 +286,39 @@ class TestPineconeDocEmbeddingsStorageProcessor:
message.metadata.user = 'test_user' message.metadata.user = 'test_user'
message.metadata.collection = 'test_collection' message.metadata.collection = 'test_collection'
chunk = ChunkEmbeddings( # Each chunk has a single vector of different dimensions
chunk=b"Document with mixed dimensions", chunk1 = ChunkEmbeddings(
vectors=[ chunk=b"Document chunk 1",
[0.1, 0.2], # 2D vector vector=[0.1, 0.2] # 2D vector
[0.3, 0.4, 0.5, 0.6], # 4D vector
[0.7, 0.8, 0.9] # 3D vector
]
) )
message.chunks = [chunk] chunk2 = ChunkEmbeddings(
chunk=b"Document chunk 2",
mock_index_2d = MagicMock() vector=[0.3, 0.4, 0.5, 0.6] # 4D vector
mock_index_4d = MagicMock() )
mock_index_3d = MagicMock() chunk3 = ChunkEmbeddings(
chunk=b"Document chunk 3",
vector=[0.7, 0.8, 0.9] # 3D vector
)
message.chunks = [chunk1, chunk2, chunk3]
mock_index = MagicMock()
def mock_index_side_effect(name): def mock_index_side_effect(name):
# All dimensions now use the same index name pattern # All dimensions now use the same index name pattern
# Different dimensions will be handled within the same index
if "test_user" in name and "test_collection" in name: if "test_user" in name and "test_collection" in name:
return mock_index_2d # Just return one mock for all return mock_index
return MagicMock() return MagicMock()
processor.pinecone.Index.side_effect = mock_index_side_effect processor.pinecone.Index.side_effect = mock_index_side_effect
processor.pinecone.has_index.return_value = True processor.pinecone.has_index.return_value = True
with patch('uuid.uuid4', side_effect=['id1', 'id2', 'id3']): with patch('uuid.uuid4', side_effect=['id1', 'id2', 'id3']):
await processor.store_document_embeddings(message) await processor.store_document_embeddings(message)
# Verify all vectors are now stored in the same index # Verify all vectors are now stored in the same index
# (Pinecone can handle mixed dimensions in the same index) # (Each chunk has a single vector, called once per chunk)
assert processor.pinecone.Index.call_count == 3 # Called once per vector assert processor.pinecone.Index.call_count == 3 # Called once per chunk
mock_index_2d.upsert.call_count == 3 # All upserts go to same index assert mock_index.upsert.call_count == 3 # All upserts go to same index
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_store_document_embeddings_empty_chunks_list(self, processor): async def test_store_document_embeddings_empty_chunks_list(self, processor):
@ -346,7 +348,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
chunk = ChunkEmbeddings( chunk = ChunkEmbeddings(
chunk=b"Document with no vectors", chunk=b"Document with no vectors",
vectors=[] vector=[]
) )
message.chunks = [chunk] message.chunks = [chunk]

View file

@ -24,16 +24,20 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
message.metadata.user = 'test_user' message.metadata.user = 'test_user'
message.metadata.collection = 'test_collection' message.metadata.collection = 'test_collection'
# Create test entity embeddings # Create test entity embeddings (each entity has a single vector)
entity1 = EntityEmbeddings( entity1 = EntityEmbeddings(
entity=Value(value="http://example.org/entity1", is_uri=True), entity=Value(value="http://example.org/entity1", is_uri=True),
vectors=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] vector=[0.1, 0.2, 0.3]
) )
entity2 = EntityEmbeddings( entity2 = EntityEmbeddings(
entity=Value(value="entity2", is_uri=False), entity=Value(value="http://example.org/entity2", is_uri=True),
vectors=[[0.7, 0.8, 0.9]] vector=[0.4, 0.5, 0.6]
) )
message.entities = [entity1, entity2] entity3 = EntityEmbeddings(
entity=Value(value="entity3", is_uri=False),
vector=[0.7, 0.8, 0.9]
)
message.entities = [entity1, entity2, entity3]
return message return message
@ -190,7 +194,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
entity = EntityEmbeddings( entity = EntityEmbeddings(
entity=Value(value="test_entity", is_uri=False), entity=Value(value="test_entity", is_uri=False),
vectors=[[0.1, 0.2, 0.3]] vector=[0.1, 0.2, 0.3]
) )
message.entities = [entity] message.entities = [entity]
@ -222,7 +226,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
entity = EntityEmbeddings( entity = EntityEmbeddings(
entity=Value(value="", is_uri=False), entity=Value(value="", is_uri=False),
vectors=[[0.1, 0.2, 0.3]] vector=[0.1, 0.2, 0.3]
) )
message.entities = [entity] message.entities = [entity]
@ -244,7 +248,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
entity = EntityEmbeddings( entity = EntityEmbeddings(
entity=Value(value=None, is_uri=False), entity=Value(value=None, is_uri=False),
vectors=[[0.1, 0.2, 0.3]] vector=[0.1, 0.2, 0.3]
) )
message.entities = [entity] message.entities = [entity]
@ -322,7 +326,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
entity = EntityEmbeddings( entity = EntityEmbeddings(
entity=Value(value="test_entity", is_uri=False), entity=Value(value="test_entity", is_uri=False),
vectors=[] vector=[]
) )
message.entities = [entity] message.entities = [entity]
@ -344,7 +348,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
entity = EntityEmbeddings( entity = EntityEmbeddings(
entity=Value(value="test_entity", is_uri=False), entity=Value(value="test_entity", is_uri=False),
vectors=[[0.1, 0.2, 0.3]] vector=[0.1, 0.2, 0.3]
) )
message.entities = [entity] message.entities = [entity]
@ -369,7 +373,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor:
entity = EntityEmbeddings( entity = EntityEmbeddings(
entity=Value(value="test_entity", is_uri=False), entity=Value(value="test_entity", is_uri=False),
vectors=[[0.1, 0.2, 0.3]] vector=[0.1, 0.2, 0.3]
) )
message.entities = [entity] message.entities = [entity]