Fixing tests

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):
+        assert response.error == error
-        """Test response translator converts Pulsar schema with chunks to dict"""
+        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 "chunk_ids" in result
-        assert result["chunks"] == []
+        assert result["chunk_ids"] == ["doc1/c1", "doc2/c2", "doc3/c3"]
 
-    def test_response_translator_from_pulsar_with_none_chunks(self):
+    def test_response_translator_from_pulsar_with_empty_chunk_ids(self):
-        """Test response translator handles None chunks"""
+        """Test response translator handles empty chunk_ids list"""
         translator = DocumentEmbeddingsResponseTranslator()
-        
+
-        response = MagicMock()
+        response = DocumentEmbeddingsResponse(
-        response.chunks = None
+            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 "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):
         """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 "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"""
@@ -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
+        # 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)