diff --git a/tests/contract/test_document_embeddings_contract.py b/tests/contract/test_document_embeddings_contract.py index e0939aaa..c35dde4b 100644 --- a/tests/contract/test_document_embeddings_contract.py +++ b/tests/contract/test_document_embeddings_contract.py @@ -25,13 +25,13 @@ class TestDocumentEmbeddingsRequestContract: user="test_user", collection="test_collection" ) - + # Verify all expected fields exist assert hasattr(request, 'vectors') assert hasattr(request, 'limit') assert hasattr(request, 'user') assert hasattr(request, 'collection') - + # Verify field values assert request.vectors == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] assert request.limit == 10 @@ -41,16 +41,16 @@ class TestDocumentEmbeddingsRequestContract: def test_request_translator_to_pulsar(self): """Test request translator converts dict to Pulsar schema""" translator = DocumentEmbeddingsRequestTranslator() - + data = { "vectors": [[0.1, 0.2], [0.3, 0.4]], "limit": 5, "user": "custom_user", "collection": "custom_collection" } - + result = translator.to_pulsar(data) - + assert isinstance(result, DocumentEmbeddingsRequest) assert result.vectors == [[0.1, 0.2], [0.3, 0.4]] assert result.limit == 5 @@ -60,14 +60,14 @@ class TestDocumentEmbeddingsRequestContract: def test_request_translator_to_pulsar_with_defaults(self): """Test request translator uses correct defaults""" translator = DocumentEmbeddingsRequestTranslator() - + data = { "vectors": [[0.1, 0.2]] # No limit, user, or collection provided } - + result = translator.to_pulsar(data) - + assert isinstance(result, DocumentEmbeddingsRequest) assert result.vectors == [[0.1, 0.2]] assert result.limit == 10 # Default @@ -77,16 +77,16 @@ class TestDocumentEmbeddingsRequestContract: def test_request_translator_from_pulsar(self): """Test request translator converts Pulsar schema to dict""" translator = DocumentEmbeddingsRequestTranslator() - + request = DocumentEmbeddingsRequest( vectors=[[0.5, 0.6]], limit=20, user="test_user", collection="test_collection" ) - + result = translator.from_pulsar(request) - + assert isinstance(result, dict) assert result["vectors"] == [[0.5, 0.6]] assert result["limit"] == 20 @@ -99,19 +99,19 @@ class TestDocumentEmbeddingsResponseContract: def test_response_schema_fields(self): """Test that DocumentEmbeddingsResponse has expected fields""" - # Create a response with chunks + # Create a response with chunk_ids response = DocumentEmbeddingsResponse( error=None, - chunks=["chunk1", "chunk2", "chunk3"] + chunk_ids=["chunk1", "chunk2", "chunk3"] ) - + # Verify all expected fields exist assert hasattr(response, 'error') - assert hasattr(response, 'chunks') - + assert hasattr(response, 'chunk_ids') + # Verify field values assert response.error is None - assert response.chunks == ["chunk1", "chunk2", "chunk3"] + assert response.chunk_ids == ["chunk1", "chunk2", "chunk3"] def test_response_schema_with_error(self): """Test response schema with error""" @@ -119,90 +119,79 @@ class TestDocumentEmbeddingsResponseContract: type="query_error", message="Database connection failed" ) - + response = DocumentEmbeddingsResponse( error=error, - chunks=None + chunk_ids=[] ) - - assert response.error == error - assert response.chunks is None - def test_response_translator_from_pulsar_with_chunks(self): - """Test response translator converts Pulsar schema with chunks to dict""" + assert response.error == error + assert response.chunk_ids == [] + + def test_response_translator_from_pulsar_with_chunk_ids(self): + """Test response translator converts Pulsar schema with chunk_ids to dict""" translator = DocumentEmbeddingsResponseTranslator() - + response = DocumentEmbeddingsResponse( error=None, - chunks=["doc1", "doc2", "doc3"] + chunk_ids=["doc1/c1", "doc2/c2", "doc3/c3"] ) - - result = translator.from_pulsar(response) - - assert isinstance(result, dict) - assert "chunks" in result - assert result["chunks"] == ["doc1", "doc2", "doc3"] - def test_response_translator_from_pulsar_with_bytes(self): - """Test response translator handles byte chunks correctly""" - translator = DocumentEmbeddingsResponseTranslator() - - response = MagicMock() - response.chunks = [b"byte_chunk1", b"byte_chunk2"] - result = translator.from_pulsar(response) - - assert isinstance(result, dict) - assert "chunks" in result - assert result["chunks"] == ["byte_chunk1", "byte_chunk2"] - def test_response_translator_from_pulsar_with_empty_chunks(self): - """Test response translator handles empty chunks list""" - translator = DocumentEmbeddingsResponseTranslator() - - response = MagicMock() - response.chunks = [] - - result = translator.from_pulsar(response) - assert isinstance(result, dict) - assert "chunks" in result - assert result["chunks"] == [] + assert "chunk_ids" in result + assert result["chunk_ids"] == ["doc1/c1", "doc2/c2", "doc3/c3"] - def test_response_translator_from_pulsar_with_none_chunks(self): - """Test response translator handles None chunks""" + def test_response_translator_from_pulsar_with_empty_chunk_ids(self): + """Test response translator handles empty chunk_ids list""" translator = DocumentEmbeddingsResponseTranslator() - - response = MagicMock() - response.chunks = None - + + response = DocumentEmbeddingsResponse( + error=None, + chunk_ids=[] + ) + result = translator.from_pulsar(response) - + assert isinstance(result, dict) - assert "chunks" not in result or result.get("chunks") is None + assert "chunk_ids" in result + assert result["chunk_ids"] == [] + + def test_response_translator_from_pulsar_with_none_chunk_ids(self): + """Test response translator handles None chunk_ids""" + translator = DocumentEmbeddingsResponseTranslator() + + response = MagicMock() + response.chunk_ids = None + + result = translator.from_pulsar(response) + + assert isinstance(result, dict) + assert "chunk_ids" not in result or result.get("chunk_ids") is None def test_response_translator_from_response_with_completion(self): """Test response translator with completion flag""" translator = DocumentEmbeddingsResponseTranslator() - + response = DocumentEmbeddingsResponse( error=None, - chunks=["chunk1", "chunk2"] + chunk_ids=["chunk1", "chunk2"] ) - + result, is_final = translator.from_response_with_completion(response) - + assert isinstance(result, dict) - assert "chunks" in result - assert result["chunks"] == ["chunk1", "chunk2"] + assert "chunk_ids" in result + assert result["chunk_ids"] == ["chunk1", "chunk2"] assert is_final is True # Document embeddings responses are always final def test_response_translator_to_pulsar_not_implemented(self): """Test that to_pulsar raises NotImplementedError for responses""" translator = DocumentEmbeddingsResponseTranslator() - + with pytest.raises(NotImplementedError): - translator.to_pulsar({"chunks": ["test"]}) + translator.to_pulsar({"chunk_ids": ["test"]}) class TestDocumentEmbeddingsMessageCompatibility: @@ -217,26 +206,26 @@ class TestDocumentEmbeddingsMessageCompatibility: "user": "test_user", "collection": "test_collection" } - + # Convert to Pulsar request req_translator = DocumentEmbeddingsRequestTranslator() pulsar_request = req_translator.to_pulsar(request_data) - + # Simulate service processing and creating response response = DocumentEmbeddingsResponse( error=None, - chunks=["relevant chunk 1", "relevant chunk 2"] + chunk_ids=["doc1/c1", "doc2/c2"] ) - + # Convert response back to dict resp_translator = DocumentEmbeddingsResponseTranslator() response_data = resp_translator.from_pulsar(response) - + # Verify data integrity assert isinstance(pulsar_request, DocumentEmbeddingsRequest) assert isinstance(response_data, dict) - assert "chunks" in response_data - assert len(response_data["chunks"]) == 2 + assert "chunk_ids" in response_data + assert len(response_data["chunk_ids"]) == 2 def test_error_response_flow(self): """Test error response flow""" @@ -245,17 +234,18 @@ class TestDocumentEmbeddingsMessageCompatibility: type="vector_db_error", message="Collection not found" ) - + response = DocumentEmbeddingsResponse( error=error, - chunks=None + chunk_ids=[] ) - + # Convert response to dict translator = DocumentEmbeddingsResponseTranslator() response_data = translator.from_pulsar(response) - + # Verify error handling assert isinstance(response_data, dict) - # The translator doesn't include error in the dict, only chunks - assert "chunks" not in response_data or response_data.get("chunks") is None \ No newline at end of file + # The translator doesn't include error in the dict, only chunk_ids + assert "chunk_ids" in response_data + assert response_data["chunk_ids"] == [] diff --git a/tests/unit/test_base/test_document_embeddings_client.py b/tests/unit/test_base/test_document_embeddings_client.py index 1c91408d..81d4a98e 100644 --- a/tests/unit/test_base/test_document_embeddings_client.py +++ b/tests/unit/test_base/test_document_embeddings_client.py @@ -22,7 +22,7 @@ class TestDocumentEmbeddingsClient(IsolatedAsyncioTestCase): client = DocumentEmbeddingsClient() mock_response = MagicMock(spec=DocumentEmbeddingsResponse) mock_response.error = None - mock_response.chunks = ["chunk1", "chunk2", "chunk3"] + mock_response.chunk_ids = ["chunk1", "chunk2", "chunk3"] # Mock the request method client.request = AsyncMock(return_value=mock_response) @@ -75,7 +75,7 @@ class TestDocumentEmbeddingsClient(IsolatedAsyncioTestCase): client = DocumentEmbeddingsClient() mock_response = MagicMock(spec=DocumentEmbeddingsResponse) mock_response.error = None - mock_response.chunks = [] + mock_response.chunk_ids = [] client.request = AsyncMock(return_value=mock_response) @@ -93,7 +93,7 @@ class TestDocumentEmbeddingsClient(IsolatedAsyncioTestCase): client = DocumentEmbeddingsClient() mock_response = MagicMock(spec=DocumentEmbeddingsResponse) mock_response.error = None - mock_response.chunks = ["test_chunk"] + mock_response.chunk_ids = ["test_chunk"] client.request = AsyncMock(return_value=mock_response) @@ -115,7 +115,7 @@ class TestDocumentEmbeddingsClient(IsolatedAsyncioTestCase): client = DocumentEmbeddingsClient() mock_response = MagicMock(spec=DocumentEmbeddingsResponse) mock_response.error = None - mock_response.chunks = ["chunk1"] + mock_response.chunk_ids = ["chunk1"] client.request = AsyncMock(return_value=mock_response) @@ -136,7 +136,7 @@ class TestDocumentEmbeddingsClient(IsolatedAsyncioTestCase): client = DocumentEmbeddingsClient() mock_response = MagicMock(spec=DocumentEmbeddingsResponse) mock_response.error = None - mock_response.chunks = ["test_chunk"] + mock_response.chunk_ids = ["test_chunk"] client.request = AsyncMock(return_value=mock_response)