Fixing tests

tests/contract/test_document_embeddings_contract.py

@@ -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"""
@@ -122,64 +122,53 @@ class TestDocumentEmbeddingsResponseContract:
 
         response = DocumentEmbeddingsResponse(
             error=error,
-            chunks=None
+            chunk_ids=[]
         )
 
         assert response.error == error
-        assert response.chunks is None
+        assert response.chunk_ids == []
 
-    def test_response_translator_from_pulsar_with_chunks(self):
+    def test_response_translator_from_pulsar_with_chunk_ids(self):
-        """Test response translator converts Pulsar schema with chunks to dict"""
+        """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 "chunk_ids" in result
-        assert result["chunks"] == ["doc1", "doc2", "doc3"]
+        assert result["chunk_ids"] == ["doc1/c1", "doc2/c2", "doc3/c3"]
 
-    def test_response_translator_from_pulsar_with_bytes(self):
+    def test_response_translator_from_pulsar_with_empty_chunk_ids(self):
-        """Test response translator handles byte chunks correctly"""
+        """Test response translator handles empty chunk_ids list"""
         translator = DocumentEmbeddingsResponseTranslator()
 
-        response = MagicMock()
+        response = DocumentEmbeddingsResponse(
-        response.chunks = [b"byte_chunk1", b"byte_chunk2"]
+            error=None,
+            chunk_ids=[]
+        )
 
         result = translator.from_pulsar(response)
 
         assert isinstance(result, dict)
-        assert "chunks" in result
+        assert "chunk_ids" in result
-        assert result["chunks"] == ["byte_chunk1", "byte_chunk2"]
+        assert result["chunk_ids"] == []
 
-    def test_response_translator_from_pulsar_with_empty_chunks(self):
+    def test_response_translator_from_pulsar_with_none_chunk_ids(self):
-        """Test response translator handles empty chunks list"""
+        """Test response translator handles None chunk_ids"""
         translator = DocumentEmbeddingsResponseTranslator()
 
         response = MagicMock()
-        response.chunks = []
+        response.chunk_ids = None
 
         result = translator.from_pulsar(response)
 
         assert isinstance(result, dict)
-        assert "chunks" in result
+        assert "chunk_ids" not in result or result.get("chunk_ids") is None
-        assert result["chunks"] == []
-
-    def test_response_translator_from_pulsar_with_none_chunks(self):
-        """Test response translator handles None chunks"""
-        translator = DocumentEmbeddingsResponseTranslator()
-        
-        response = MagicMock()
-        response.chunks = None
-        
-        result = translator.from_pulsar(response)
-        
-        assert isinstance(result, dict)
-        assert "chunks" not in result or result.get("chunks") is None
 
     def test_response_translator_from_response_with_completion(self):
         """Test response translator with completion flag"""
@@ -187,14 +176,14 @@ class TestDocumentEmbeddingsResponseContract:
 
         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 "chunk_ids" in result
-        assert result["chunks"] == ["chunk1", "chunk2"]
+        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):
@@ -202,7 +191,7 @@ class TestDocumentEmbeddingsResponseContract:
         translator = DocumentEmbeddingsResponseTranslator()
 
         with pytest.raises(NotImplementedError):
-            translator.to_pulsar({"chunks": ["test"]})
+            translator.to_pulsar({"chunk_ids": ["test"]})
 
 
 class TestDocumentEmbeddingsMessageCompatibility:
@@ -225,7 +214,7 @@ class TestDocumentEmbeddingsMessageCompatibility:
         # 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
@@ -235,8 +224,8 @@ class TestDocumentEmbeddingsMessageCompatibility:
         # Verify data integrity
         assert isinstance(pulsar_request, DocumentEmbeddingsRequest)
         assert isinstance(response_data, dict)
-        assert "chunks" in response_data
+        assert "chunk_ids" in response_data
-        assert len(response_data["chunks"]) == 2
+        assert len(response_data["chunk_ids"]) == 2
 
     def test_error_response_flow(self):
         """Test error response flow"""
@@ -248,7 +237,7 @@ class TestDocumentEmbeddingsMessageCompatibility:
 
         response = DocumentEmbeddingsResponse(
             error=error,
-            chunks=None
+            chunk_ids=[]
         )
 
         # Convert response to dict
@@ -257,5 +246,6 @@ class TestDocumentEmbeddingsMessageCompatibility:
 
         # Verify error handling
         assert isinstance(response_data, dict)
-        # The translator doesn't include error in the dict, only chunks
+        # The translator doesn't include error in the dict, only chunk_ids
-        assert "chunks" not in response_data or response_data.get("chunks") is None
+        assert "chunk_ids" in response_data
+        assert response_data["chunk_ids"] == []

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)