mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-14 15:52:11 +02:00
Fixing tests
This commit is contained in:
parent
f89e002614
commit
8d25f4d4a8
3 changed files with 39 additions and 33 deletions
|
|
@ -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]
|
||||||
)
|
)
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -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]
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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]
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue