From 2d47102d707cb7cc8e92a90e7b6e5f936f07ed3a Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Sat, 7 Mar 2026 23:32:25 +0000 Subject: [PATCH] Fixing tests --- .../test_doc_embeddings_milvus_query.py | 52 +++--- .../test_doc_embeddings_qdrant_query.py | 34 ++-- .../test_doc_embeddings_milvus_storage.py | 74 ++++---- .../test_doc_embeddings_qdrant_storage.py | 159 +++++++----------- 4 files changed, 141 insertions(+), 178 deletions(-) diff --git a/tests/unit/test_query/test_doc_embeddings_milvus_query.py b/tests/unit/test_query/test_doc_embeddings_milvus_query.py index 622529e5..b4c954d8 100644 --- a/tests/unit/test_query/test_doc_embeddings_milvus_query.py +++ b/tests/unit/test_query/test_doc_embeddings_milvus_query.py @@ -77,9 +77,9 @@ class TestMilvusDocEmbeddingsQueryProcessor: # Mock search results mock_results = [ - {"entity": {"doc": "First document chunk"}}, - {"entity": {"doc": "Second document chunk"}}, - {"entity": {"doc": "Third document chunk"}}, + {"entity": {"chunk_id": "First document chunk"}}, + {"entity": {"chunk_id": "Second document chunk"}}, + {"entity": {"chunk_id": "Third document chunk"}}, ] processor.vecstore.search.return_value = mock_results @@ -108,11 +108,11 @@ class TestMilvusDocEmbeddingsQueryProcessor: # Mock search results - different results for each vector mock_results_1 = [ - {"entity": {"doc": "Document from first vector"}}, - {"entity": {"doc": "Another doc from first vector"}}, + {"entity": {"chunk_id": "Document from first vector"}}, + {"entity": {"chunk_id": "Another doc from first vector"}}, ] mock_results_2 = [ - {"entity": {"doc": "Document from second vector"}}, + {"entity": {"chunk_id": "Document from second vector"}}, ] processor.vecstore.search.side_effect = [mock_results_1, mock_results_2] @@ -147,10 +147,10 @@ class TestMilvusDocEmbeddingsQueryProcessor: # Mock search results - more results than limit mock_results = [ - {"entity": {"doc": "Document 1"}}, - {"entity": {"doc": "Document 2"}}, - {"entity": {"doc": "Document 3"}}, - {"entity": {"doc": "Document 4"}}, + {"entity": {"chunk_id": "Document 1"}}, + {"entity": {"chunk_id": "Document 2"}}, + {"entity": {"chunk_id": "Document 3"}}, + {"entity": {"chunk_id": "Document 4"}}, ] processor.vecstore.search.return_value = mock_results @@ -217,9 +217,9 @@ class TestMilvusDocEmbeddingsQueryProcessor: # Mock search results with Unicode content mock_results = [ - {"entity": {"doc": "Document with Unicode: éñ中文🚀"}}, - {"entity": {"doc": "Regular ASCII document"}}, - {"entity": {"doc": "Document with émojis: 😀🎉"}}, + {"entity": {"chunk_id": "Document with Unicode: éñ中文🚀"}}, + {"entity": {"chunk_id": "Regular ASCII document"}}, + {"entity": {"chunk_id": "Document with émojis: 😀🎉"}}, ] processor.vecstore.search.return_value = mock_results @@ -244,8 +244,8 @@ class TestMilvusDocEmbeddingsQueryProcessor: # Mock search results with large content large_doc = "A" * 10000 # 10KB of content mock_results = [ - {"entity": {"doc": large_doc}}, - {"entity": {"doc": "Small document"}}, + {"entity": {"chunk_id": large_doc}}, + {"entity": {"chunk_id": "Small document"}}, ] processor.vecstore.search.return_value = mock_results @@ -268,9 +268,9 @@ class TestMilvusDocEmbeddingsQueryProcessor: # Mock search results with special characters mock_results = [ - {"entity": {"doc": "Document with \"quotes\" and 'apostrophes'"}}, - {"entity": {"doc": "Document with\nnewlines\tand\ttabs"}}, - {"entity": {"doc": "Document with special chars: @#$%^&*()"}}, + {"entity": {"chunk_id": "Document with \"quotes\" and 'apostrophes'"}}, + {"entity": {"chunk_id": "Document with\nnewlines\tand\ttabs"}}, + {"entity": {"chunk_id": "Document with special chars: @#$%^&*()"}}, ] processor.vecstore.search.return_value = mock_results @@ -350,9 +350,9 @@ class TestMilvusDocEmbeddingsQueryProcessor: ) # Mock search results for each vector - mock_results_1 = [{"entity": {"doc": "Document from 2D vector"}}] - mock_results_2 = [{"entity": {"doc": "Document from 4D vector"}}] - mock_results_3 = [{"entity": {"doc": "Document from 3D vector"}}] + mock_results_1 = [{"entity": {"chunk_id": "Document from 2D vector"}}] + mock_results_2 = [{"entity": {"chunk_id": "Document from 4D vector"}}] + mock_results_3 = [{"entity": {"chunk_id": "Document from 3D vector"}}] processor.vecstore.search.side_effect = [mock_results_1, mock_results_2, mock_results_3] result = await processor.query_document_embeddings(query) @@ -378,12 +378,12 @@ class TestMilvusDocEmbeddingsQueryProcessor: # Mock search results with duplicates across vectors mock_results_1 = [ - {"entity": {"doc": "Document A"}}, - {"entity": {"doc": "Document B"}}, + {"entity": {"chunk_id": "Document A"}}, + {"entity": {"chunk_id": "Document B"}}, ] mock_results_2 = [ - {"entity": {"doc": "Document B"}}, # Duplicate - {"entity": {"doc": "Document C"}}, + {"entity": {"chunk_id": "Document B"}}, # Duplicate + {"entity": {"chunk_id": "Document C"}}, ] processor.vecstore.search.side_effect = [mock_results_1, mock_results_2] @@ -458,5 +458,5 @@ class TestMilvusDocEmbeddingsQueryProcessor: mock_launch.assert_called_once_with( default_ident, - "\nDocument embeddings query service. Input is vector, output is an array\nof chunks\n" + "\nDocument embeddings query service. Input is vector, output is an array\nof chunk_ids\n" ) \ No newline at end of file diff --git a/tests/unit/test_query/test_doc_embeddings_qdrant_query.py b/tests/unit/test_query/test_doc_embeddings_qdrant_query.py index ad337f73..204c4cc3 100644 --- a/tests/unit/test_query/test_doc_embeddings_qdrant_query.py +++ b/tests/unit/test_query/test_doc_embeddings_qdrant_query.py @@ -77,9 +77,9 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase): # Mock query response mock_point1 = MagicMock() - mock_point1.payload = {'doc': 'first document chunk'} + mock_point1.payload = {'chunk_id': 'first document chunk'} mock_point2 = MagicMock() - mock_point2.payload = {'doc': 'second document chunk'} + mock_point2.payload = {'chunk_id': 'second document chunk'} mock_response = MagicMock() mock_response.points = [mock_point1, mock_point2] @@ -132,11 +132,11 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase): # Mock query responses for different vectors mock_point1 = MagicMock() - mock_point1.payload = {'doc': 'document from vector 1'} + mock_point1.payload = {'chunk_id': 'document from vector 1'} mock_point2 = MagicMock() - mock_point2.payload = {'doc': 'document from vector 2'} + mock_point2.payload = {'chunk_id': 'document from vector 2'} mock_point3 = MagicMock() - mock_point3.payload = {'doc': 'another document from vector 2'} + mock_point3.payload = {'chunk_id': 'another document from vector 2'} mock_response1 = MagicMock() mock_response1.points = [mock_point1] @@ -192,7 +192,7 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase): mock_points = [] for i in range(10): mock_point = MagicMock() - mock_point.payload = {'doc': f'document chunk {i}'} + mock_point.payload = {'chunk_id': f'document chunk {i}'} mock_points.append(mock_point) mock_response = MagicMock() @@ -270,9 +270,9 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase): # Mock query responses mock_point1 = MagicMock() - mock_point1.payload = {'doc': 'document from 2D vector'} + mock_point1.payload = {'chunk_id': 'document from 2D vector'} mock_point2 = MagicMock() - mock_point2.payload = {'doc': 'document from 3D vector'} + mock_point2.payload = {'chunk_id': 'document from 3D vector'} mock_response1 = MagicMock() mock_response1.points = [mock_point1] @@ -326,9 +326,9 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase): # Mock query response with UTF-8 content mock_point1 = MagicMock() - mock_point1.payload = {'doc': 'Document with UTF-8: café, naïve, résumé'} + mock_point1.payload = {'chunk_id': 'Document with UTF-8: café, naïve, résumé'} mock_point2 = MagicMock() - mock_point2.payload = {'doc': 'Chinese text: 你好世界'} + mock_point2.payload = {'chunk_id': 'Chinese text: 你好世界'} mock_response = MagicMock() mock_response.points = [mock_point1, mock_point2] @@ -399,7 +399,7 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase): # Mock query response mock_point = MagicMock() - mock_point.payload = {'doc': 'document chunk'} + mock_point.payload = {'chunk_id': 'document chunk'} mock_response = MagicMock() mock_response.points = [mock_point] mock_qdrant_instance.query_points.return_value = mock_response @@ -442,9 +442,9 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase): # Mock query response with fewer results than limit mock_point1 = MagicMock() - mock_point1.payload = {'doc': 'document 1'} + mock_point1.payload = {'chunk_id': 'document 1'} mock_point2 = MagicMock() - mock_point2.payload = {'doc': 'document 2'} + mock_point2.payload = {'chunk_id': 'document 2'} mock_response = MagicMock() mock_response.points = [mock_point1, mock_point2] @@ -487,11 +487,11 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase): mock_qdrant_instance = MagicMock() mock_qdrant_client.return_value = mock_qdrant_instance - # Mock query response with missing 'doc' key + # Mock query response with missing 'chunk_id' key mock_point1 = MagicMock() - mock_point1.payload = {'doc': 'valid document'} + mock_point1.payload = {'chunk_id': 'valid document'} mock_point2 = MagicMock() - mock_point2.payload = {} # Missing 'doc' key + mock_point2.payload = {} # Missing 'chunk_id' key mock_point3 = MagicMock() mock_point3.payload = {'other_key': 'invalid'} # Wrong key @@ -514,7 +514,7 @@ class TestQdrantDocEmbeddingsQuery(IsolatedAsyncioTestCase): mock_message.collection = 'payload_collection' # Act & Assert - # This should raise a KeyError when trying to access payload['doc'] + # This should raise a KeyError when trying to access payload['chunk_id'] with pytest.raises(KeyError): await processor.query_document_embeddings(mock_message) 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 d957d711..8da88fcb 100644 --- a/tests/unit/test_storage/test_doc_embeddings_milvus_storage.py +++ b/tests/unit/test_storage/test_doc_embeddings_milvus_storage.py @@ -22,11 +22,11 @@ class TestMilvusDocEmbeddingsStorageProcessor: # Create test document embeddings chunk1 = ChunkEmbeddings( - chunk=b"This is the first document chunk", + chunk_id="This is the first document chunk", vectors=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] ) chunk2 = ChunkEmbeddings( - chunk=b"This is the second document chunk", + chunk_id="This is the second document chunk", vectors=[[0.7, 0.8, 0.9]] ) message.chunks = [chunk1, chunk2] @@ -84,7 +84,7 @@ class TestMilvusDocEmbeddingsStorageProcessor: message.metadata.collection = 'test_collection' chunk = ChunkEmbeddings( - chunk=b"Test document content", + chunk_id="Test document content", vectors=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] ) message.chunks = [chunk] @@ -136,7 +136,7 @@ class TestMilvusDocEmbeddingsStorageProcessor: message.metadata.collection = 'test_collection' chunk = ChunkEmbeddings( - chunk=b"", + chunk_id="", vectors=[[0.1, 0.2, 0.3]] ) message.chunks = [chunk] @@ -148,21 +148,21 @@ class TestMilvusDocEmbeddingsStorageProcessor: @pytest.mark.asyncio async def test_store_document_embeddings_none_chunk(self, processor): - """Test storing document embeddings with None chunk (should be skipped)""" + """Test storing document embeddings with None chunk_id (should be skipped)""" message = MagicMock() message.metadata = MagicMock() message.metadata.user = 'test_user' message.metadata.collection = 'test_collection' - + chunk = ChunkEmbeddings( - chunk=None, + chunk_id=None, vectors=[[0.1, 0.2, 0.3]] ) message.chunks = [chunk] - + await processor.store_document_embeddings(message) - - # Verify no insert was called for None chunk + + # Verify no insert was called for None chunk_id processor.vecstore.insert.assert_not_called() @pytest.mark.asyncio @@ -174,15 +174,15 @@ class TestMilvusDocEmbeddingsStorageProcessor: message.metadata.collection = 'test_collection' valid_chunk = ChunkEmbeddings( - chunk=b"Valid document content", + chunk_id="Valid document content", vectors=[[0.1, 0.2, 0.3]] ) empty_chunk = ChunkEmbeddings( - chunk=b"", + chunk_id="", vectors=[[0.4, 0.5, 0.6]] ) none_chunk = ChunkEmbeddings( - chunk=None, + chunk_id=None, vectors=[[0.7, 0.8, 0.9]] ) message.chunks = [valid_chunk, empty_chunk, none_chunk] @@ -217,7 +217,7 @@ class TestMilvusDocEmbeddingsStorageProcessor: message.metadata.collection = 'test_collection' chunk = ChunkEmbeddings( - chunk=b"Document with no vectors", + chunk_id="Document with no vectors", vectors=[] ) message.chunks = [chunk] @@ -236,7 +236,7 @@ class TestMilvusDocEmbeddingsStorageProcessor: message.metadata.collection = 'test_collection' chunk = ChunkEmbeddings( - chunk=b"Document with mixed dimensions", + chunk_id="Document with mixed dimensions", vectors=[ [0.1, 0.2], # 2D vector [0.3, 0.4, 0.5, 0.6], # 4D vector @@ -264,46 +264,46 @@ class TestMilvusDocEmbeddingsStorageProcessor: @pytest.mark.asyncio async def test_store_document_embeddings_unicode_content(self, processor): - """Test storing document embeddings with Unicode content""" + """Test storing document embeddings with Unicode content in chunk_id""" message = MagicMock() message.metadata = MagicMock() message.metadata.user = 'test_user' message.metadata.collection = 'test_collection' - + chunk = ChunkEmbeddings( - chunk="Document with Unicode: éñ中文🚀".encode('utf-8'), + chunk_id="chunk/doc/unicode-éñ中文🚀", vectors=[[0.1, 0.2, 0.3]] ) message.chunks = [chunk] - + await processor.store_document_embeddings(message) - - # Verify Unicode content was properly decoded and inserted with user/collection parameters + + # Verify Unicode chunk_id was stored correctly with user/collection parameters processor.vecstore.insert.assert_called_once_with( - [0.1, 0.2, 0.3], "Document with Unicode: éñ中文🚀", 'test_user', 'test_collection' + [0.1, 0.2, 0.3], "chunk/doc/unicode-éñ中文🚀", 'test_user', 'test_collection' ) @pytest.mark.asyncio - async def test_store_document_embeddings_large_chunks(self, processor): - """Test storing document embeddings with large document chunks""" + async def test_store_document_embeddings_large_chunk_id(self, processor): + """Test storing document embeddings with long chunk_id""" message = MagicMock() message.metadata = MagicMock() message.metadata.user = 'test_user' message.metadata.collection = 'test_collection' - - # Create a large document chunk - large_content = "A" * 10000 # 10KB of content + + # Create a long chunk_id + long_chunk_id = "chunk/doc/" + "a" * 200 chunk = ChunkEmbeddings( - chunk=large_content.encode('utf-8'), + chunk_id=long_chunk_id, vectors=[[0.1, 0.2, 0.3]] ) message.chunks = [chunk] - + await processor.store_document_embeddings(message) - - # Verify large content was inserted with user/collection parameters + + # Verify long chunk_id was inserted with user/collection parameters processor.vecstore.insert.assert_called_once_with( - [0.1, 0.2, 0.3], large_content, 'test_user', 'test_collection' + [0.1, 0.2, 0.3], long_chunk_id, 'test_user', 'test_collection' ) @pytest.mark.asyncio @@ -315,7 +315,7 @@ class TestMilvusDocEmbeddingsStorageProcessor: message.metadata.collection = 'test_collection' chunk = ChunkEmbeddings( - chunk=b" \n\t ", + chunk_id=" \n\t ", vectors=[[0.1, 0.2, 0.3]] ) message.chunks = [chunk] @@ -346,7 +346,7 @@ class TestMilvusDocEmbeddingsStorageProcessor: message.metadata.collection = collection chunk = ChunkEmbeddings( - chunk=b"Test content", + chunk_id="Test content", vectors=[[0.1, 0.2, 0.3]] ) message.chunks = [chunk] @@ -367,7 +367,7 @@ class TestMilvusDocEmbeddingsStorageProcessor: message1.metadata.user = 'user1' message1.metadata.collection = 'collection1' chunk1 = ChunkEmbeddings( - chunk=b"User1 content", + chunk_id="User1 content", vectors=[[0.1, 0.2, 0.3]] ) message1.chunks = [chunk1] @@ -378,7 +378,7 @@ class TestMilvusDocEmbeddingsStorageProcessor: message2.metadata.user = 'user2' message2.metadata.collection = 'collection2' chunk2 = ChunkEmbeddings( - chunk=b"User2 content", + chunk_id="User2 content", vectors=[[0.4, 0.5, 0.6]] ) message2.chunks = [chunk2] @@ -409,7 +409,7 @@ class TestMilvusDocEmbeddingsStorageProcessor: message.metadata.collection = 'test-collection.v1' # Collection with special chars chunk = ChunkEmbeddings( - chunk=b"Special chars test", + chunk_id="Special chars test", vectors=[[0.1, 0.2, 0.3]] ) message.chunks = [chunk] diff --git a/tests/unit/test_storage/test_doc_embeddings_qdrant_storage.py b/tests/unit/test_storage/test_doc_embeddings_qdrant_storage.py index fc839482..e3498ce9 100644 --- a/tests/unit/test_storage/test_doc_embeddings_qdrant_storage.py +++ b/tests/unit/test_storage/test_doc_embeddings_qdrant_storage.py @@ -20,7 +20,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): # Arrange mock_qdrant_instance = MagicMock() mock_qdrant_client.return_value = mock_qdrant_instance - + config = { 'store_uri': 'http://localhost:6333', 'api_key': 'test-api-key', @@ -34,7 +34,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): # Assert # Verify QdrantClient was created with correct parameters mock_qdrant_client.assert_called_once_with(url='http://localhost:6333', api_key='test-api-key') - + # Verify processor attributes assert hasattr(processor, 'qdrant') assert processor.qdrant == mock_qdrant_instance @@ -45,7 +45,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): # Arrange mock_qdrant_instance = MagicMock() mock_qdrant_client.return_value = mock_qdrant_instance - + config = { 'taskgroup': AsyncMock(), 'id': 'test-doc-qdrant-processor' @@ -69,7 +69,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): 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-123') - + config = { 'store_uri': 'http://localhost:6333', 'api_key': 'test-api-key', @@ -86,13 +86,13 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): mock_message = MagicMock() mock_message.metadata.user = 'test_user' mock_message.metadata.collection = 'test_collection' - + mock_chunk = MagicMock() - mock_chunk.chunk.decode.return_value = 'test document chunk' + mock_chunk.chunk_id = 'doc/c1' # chunk_id instead of chunk bytes mock_chunk.vectors = [[0.1, 0.2, 0.3]] # Single vector with 3 dimensions - + mock_message.chunks = [mock_chunk] - + # Act await processor.store_document_embeddings(mock_message) @@ -100,18 +100,18 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): # Verify collection existence was checked (with dimension suffix) expected_collection = 'd_test_user_test_collection_3' # 3 dimensions in vector [0.1, 0.2, 0.3] mock_qdrant_instance.collection_exists.assert_called_once_with(expected_collection) - + # Verify upsert was called mock_qdrant_instance.upsert.assert_called_once() - + # Verify upsert parameters upsert_call_args = mock_qdrant_instance.upsert.call_args assert upsert_call_args[1]['collection_name'] == 'd_test_user_test_collection_3' assert len(upsert_call_args[1]['points']) == 1 - + point = upsert_call_args[1]['points'][0] assert point.vector == [0.1, 0.2, 0.3] - assert point.payload['doc'] == 'test document chunk' + assert point.payload['chunk_id'] == 'doc/c1' @patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid') @@ -123,7 +123,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): 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 = { 'store_uri': 'http://localhost:6333', 'api_key': 'test-api-key', @@ -140,38 +140,38 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): mock_message = MagicMock() mock_message.metadata.user = 'multi_user' mock_message.metadata.collection = 'multi_collection' - + mock_chunk1 = MagicMock() - mock_chunk1.chunk.decode.return_value = 'first document chunk' + mock_chunk1.chunk_id = 'doc/c1' mock_chunk1.vectors = [[0.1, 0.2]] - + mock_chunk2 = MagicMock() - mock_chunk2.chunk.decode.return_value = 'second document chunk' + mock_chunk2.chunk_id = 'doc/c2' mock_chunk2.vectors = [[0.3, 0.4]] - + mock_message.chunks = [mock_chunk1, mock_chunk2] - + # Act await processor.store_document_embeddings(mock_message) # Assert # Should be called twice (once per chunk) assert mock_qdrant_instance.upsert.call_count == 2 - + # Verify both chunks were processed upsert_calls = mock_qdrant_instance.upsert.call_args_list - + # First chunk first_call = upsert_calls[0] first_point = first_call[1]['points'][0] assert first_point.vector == [0.1, 0.2] - assert first_point.payload['doc'] == 'first document chunk' - + assert first_point.payload['chunk_id'] == 'doc/c1' + # Second chunk second_call = upsert_calls[1] second_point = second_call[1]['points'][0] assert second_point.vector == [0.3, 0.4] - assert second_point.payload['doc'] == 'second document chunk' + assert second_point.payload['chunk_id'] == 'doc/c2' @patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid') @@ -183,7 +183,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): 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 = { 'store_uri': 'http://localhost:6333', 'api_key': 'test-api-key', @@ -200,41 +200,41 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): mock_message = MagicMock() mock_message.metadata.user = 'vector_user' mock_message.metadata.collection = 'vector_collection' - + mock_chunk = MagicMock() - mock_chunk.chunk.decode.return_value = 'multi-vector document chunk' + mock_chunk.chunk_id = 'doc/multi-vector' mock_chunk.vectors = [ [0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9] ] - + mock_message.chunks = [mock_chunk] - + # Act await processor.store_document_embeddings(mock_message) # Assert # Should be called 3 times (once per vector) assert mock_qdrant_instance.upsert.call_count == 3 - + # Verify all vectors 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.4, 0.5, 0.6], [0.7, 0.8, 0.9] ] - + for i, call in enumerate(upsert_calls): point = call[1]['points'][0] assert point.vector == expected_vectors[i] - assert point.payload['doc'] == 'multi-vector document chunk' + assert point.payload['chunk_id'] == 'doc/multi-vector' @patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient') - async def test_store_document_embeddings_empty_chunk(self, mock_qdrant_client): - """Test storing document embeddings skips empty chunks""" + async def test_store_document_embeddings_empty_chunk_id(self, mock_qdrant_client): + """Test storing document embeddings skips empty chunk_ids""" # Arrange mock_qdrant_instance = MagicMock() mock_qdrant_instance.collection_exists.return_value = True # Collection exists @@ -249,13 +249,13 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): processor = Processor(**config) - # Create mock message with empty chunk + # Create mock message with empty chunk_id mock_message = MagicMock() mock_message.metadata.user = 'empty_user' mock_message.metadata.collection = 'empty_collection' mock_chunk_empty = MagicMock() - mock_chunk_empty.chunk.decode.return_value = "" # Empty string + mock_chunk_empty.chunk_id = "" # Empty chunk_id mock_chunk_empty.vectors = [[0.1, 0.2]] mock_message.chunks = [mock_chunk_empty] @@ -264,9 +264,9 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): await processor.store_document_embeddings(mock_message) # Assert - # Should not call upsert for empty chunks + # Should not call upsert for empty chunk_ids mock_qdrant_instance.upsert.assert_not_called() - # collection_exists should NOT be called since we return early for empty chunks + # collection_exists should NOT be called since we return early for empty chunk_ids mock_qdrant_instance.collection_exists.assert_not_called() @patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient') @@ -298,7 +298,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): mock_message.metadata.collection = 'new_collection' mock_chunk = MagicMock() - mock_chunk.chunk.decode.return_value = 'test chunk' + mock_chunk.chunk_id = 'doc/test-chunk' mock_chunk.vectors = [[0.1, 0.2, 0.3, 0.4, 0.5]] # 5 dimensions mock_message.chunks = [mock_chunk] @@ -350,7 +350,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): mock_message.metadata.collection = 'error_collection' mock_chunk = MagicMock() - mock_chunk.chunk.decode.return_value = 'test chunk' + mock_chunk.chunk_id = 'doc/test-chunk' mock_chunk.vectors = [[0.1, 0.2]] mock_message.chunks = [mock_chunk] @@ -388,7 +388,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): mock_message1.metadata.collection = 'cache_collection' mock_chunk1 = MagicMock() - mock_chunk1.chunk.decode.return_value = 'first chunk' + mock_chunk1.chunk_id = 'doc/c1' mock_chunk1.vectors = [[0.1, 0.2, 0.3]] mock_message1.chunks = [mock_chunk1] @@ -406,7 +406,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): mock_message2.metadata.collection = 'cache_collection' mock_chunk2 = MagicMock() - mock_chunk2.chunk.decode.return_value = 'second chunk' + mock_chunk2.chunk_id = 'doc/c2' mock_chunk2.vectors = [[0.4, 0.5, 0.6]] # Same dimension (3) mock_message2.chunks = [mock_chunk2] @@ -452,7 +452,7 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): mock_message.metadata.collection = 'dim_collection' mock_chunk = MagicMock() - mock_chunk.chunk.decode.return_value = 'dimension test chunk' + mock_chunk.chunk_id = 'doc/dim-test' mock_chunk.vectors = [ [0.1, 0.2], # 2 dimensions [0.3, 0.4, 0.5] # 3 dimensions @@ -485,28 +485,28 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): # Arrange mock_qdrant_client.return_value = MagicMock() mock_parser = MagicMock() - + # Act with patch('trustgraph.base.DocumentEmbeddingsStoreService.add_args') as mock_parent_add_args: Processor.add_args(mock_parser) # Assert mock_parent_add_args.assert_called_once_with(mock_parser) - + # Verify processor-specific arguments were added assert mock_parser.add_argument.call_count >= 2 # At least store-uri and api-key @patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient') @patch('trustgraph.storage.doc_embeddings.qdrant.write.uuid') - async def test_utf8_decoding_handling(self, mock_uuid, mock_qdrant_client): - """Test proper UTF-8 decoding of chunk text""" + async def test_chunk_id_with_special_characters(self, mock_uuid, mock_qdrant_client): + """Test storing chunk_id with special characters (URIs)""" # 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 = MagicMock() mock_uuid.uuid4.return_value.__str__ = MagicMock(return_value='test-uuid') - + config = { 'store_uri': 'http://localhost:6333', 'api_key': 'test-api-key', @@ -517,65 +517,28 @@ class TestQdrantDocEmbeddingsStorage(IsolatedAsyncioTestCase): processor = Processor(**config) # Add collection to known_collections (simulates config push) - processor.known_collections[('utf8_user', 'utf8_collection')] = {} + processor.known_collections[('uri_user', 'uri_collection')] = {} - # Create mock message with UTF-8 encoded text + # Create mock message with URI-style chunk_id mock_message = MagicMock() - mock_message.metadata.user = 'utf8_user' - mock_message.metadata.collection = 'utf8_collection' - + mock_message.metadata.user = 'uri_user' + mock_message.metadata.collection = 'uri_collection' + mock_chunk = MagicMock() - mock_chunk.chunk.decode.return_value = 'UTF-8 text with special chars: café, naïve, résumé' + mock_chunk.chunk_id = 'https://trustgraph.ai/doc/my-document/p1/c3' mock_chunk.vectors = [[0.1, 0.2]] - + mock_message.chunks = [mock_chunk] - + # Act await processor.store_document_embeddings(mock_message) # Assert - # Verify chunk.decode was called with 'utf-8' - mock_chunk.chunk.decode.assert_called_with('utf-8') - - # Verify the decoded text was stored in payload + # Verify the chunk_id was stored correctly upsert_call_args = mock_qdrant_instance.upsert.call_args point = upsert_call_args[1]['points'][0] - assert point.payload['doc'] == 'UTF-8 text with special chars: café, naïve, résumé' - - @patch('trustgraph.storage.doc_embeddings.qdrant.write.QdrantClient') - async def test_chunk_decode_exception_handling(self, mock_qdrant_client): - """Test handling of chunk decode exceptions""" - # Arrange - mock_qdrant_instance = MagicMock() - mock_qdrant_client.return_value = mock_qdrant_instance - - config = { - 'store_uri': 'http://localhost:6333', - 'api_key': 'test-api-key', - 'taskgroup': AsyncMock(), - 'id': 'test-doc-qdrant-processor' - } - - processor = Processor(**config) - - # Add collection to known_collections (simulates config push) - processor.known_collections[('decode_user', 'decode_collection')] = {} - - # Create mock message with decode error - mock_message = MagicMock() - mock_message.metadata.user = 'decode_user' - mock_message.metadata.collection = 'decode_collection' - - mock_chunk = MagicMock() - mock_chunk.chunk.decode.side_effect = UnicodeDecodeError('utf-8', b'', 0, 1, 'invalid start byte') - mock_chunk.vectors = [[0.1, 0.2]] - - mock_message.chunks = [mock_chunk] - - # Act & Assert - with pytest.raises(UnicodeDecodeError): - await processor.store_document_embeddings(mock_message) + assert point.payload['chunk_id'] == 'https://trustgraph.ai/doc/my-document/p1/c3' if __name__ == '__main__': - pytest.main([__file__]) \ No newline at end of file + pytest.main([__file__])