Tests for docling-serve chunker

tests/test_chunker.py

@@ -1,9 +1,12 @@
+from unittest.mock import Mock, patch
+
 import pytest
 from datasets import Dataset
 from transformers import AutoTokenizer
 
 from haiku.rag.chunkers import get_chunker
 from haiku.rag.chunkers.docling_local import DoclingLocalChunker
+from haiku.rag.chunkers.docling_serve import DoclingServeChunker
 from haiku.rag.config import AppConfig, Config
 from haiku.rag.converters import get_converter
 
@@ -137,3 +140,171 @@ async def test_local_chunker_markdown_tables():
     table_content = [chunk for chunk in chunks_narrative if "Value" in chunk][0]
     # Narrative format uses commas, not pipes for table structure
     assert "," in table_content and "|" not in table_content
+
+
+def test_get_chunker_docling_serve():
+    """Test factory returns DoclingServeChunker for docling-serve."""
+    config = AppConfig()
+    config.processing.chunker = "docling-serve"
+    chunker = get_chunker(config)
+    assert isinstance(chunker, DoclingServeChunker)
+
+
+class TestDoclingServeChunker:
+    """Tests for DoclingServeChunker (mocked)."""
+
+    @pytest.fixture
+    def config(self):
+        """Create test configuration."""
+        config = AppConfig()
+        config.providers.docling_serve.base_url = "http://localhost:5001"
+        config.providers.docling_serve.api_key = ""
+        config.providers.docling_serve.timeout = 300
+        config.processing.chunk_size = 256
+        config.processing.chunking_tokenizer = "Qwen/Qwen3-Embedding-0.6B"
+        return config
+
+    @pytest.fixture
+    def chunker(self, config):
+        """Create DoclingServeChunker instance."""
+        return DoclingServeChunker(config)
+
+    @pytest.mark.asyncio
+    @patch("haiku.rag.chunkers.docling_serve.requests.post")
+    async def test_chunk_success(self, mock_post, chunker):
+        """Test successful chunking via docling-serve."""
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "chunks": [
+                {"text": "Chunk 1", "chunk_index": 0},
+                {"text": "Chunk 2", "chunk_index": 1},
+            ]
+        }
+        mock_post.return_value = mock_response
+
+        # Create a simple document
+        converter = get_converter(Config)
+        doc = converter.convert_text("# Test\n\nContent", name="test.md")
+
+        chunks = await chunker.chunk(doc)
+        assert len(chunks) == 2
+        assert chunks[0] == "Chunk 1"
+        assert chunks[1] == "Chunk 2"
+        mock_post.assert_called_once()
+
+    @pytest.mark.asyncio
+    @patch("haiku.rag.chunkers.docling_serve.requests.post")
+    async def test_chunk_with_api_key(self, mock_post, config):
+        """Test that API key is included in request headers."""
+        config.providers.docling_serve.api_key = "test-key"
+        chunker = DoclingServeChunker(config)
+
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "chunks": [{"text": "Chunk 1", "chunk_index": 0}]
+        }
+        mock_post.return_value = mock_response
+
+        converter = get_converter(Config)
+        doc = converter.convert_text("# Test", name="test.md")
+        await chunker.chunk(doc)
+
+        call_kwargs = mock_post.call_args.kwargs
+        assert "headers" in call_kwargs
+        assert call_kwargs["headers"]["X-Api-Key"] == "test-key"
+
+    @pytest.mark.asyncio
+    @patch("haiku.rag.chunkers.docling_serve.requests.post")
+    async def test_chunk_hierarchical_endpoint(self, mock_post, config):
+        """Test that hierarchical chunker uses correct endpoint."""
+        config.processing.chunker_type = "hierarchical"
+        chunker = DoclingServeChunker(config)
+
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "chunks": [{"text": "Chunk 1", "chunk_index": 0}]
+        }
+        mock_post.return_value = mock_response
+
+        converter = get_converter(Config)
+        doc = converter.convert_text("# Test", name="test.md")
+        await chunker.chunk(doc)
+
+        call_args = mock_post.call_args
+        assert "/v1/chunk/hierarchical/file" in call_args[0][0]
+
+    @pytest.mark.asyncio
+    @patch("haiku.rag.chunkers.docling_serve.requests.post")
+    async def test_chunk_passes_config_parameters(self, mock_post, config):
+        """Test that all config parameters are passed to API."""
+        config.processing.chunk_size = 512
+        config.processing.chunking_merge_peers = False
+        config.processing.chunking_use_markdown_tables = True
+        chunker = DoclingServeChunker(config)
+
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {
+            "chunks": [{"text": "Chunk 1", "chunk_index": 0}]
+        }
+        mock_post.return_value = mock_response
+
+        converter = get_converter(Config)
+        doc = converter.convert_text("# Test", name="test.md")
+        await chunker.chunk(doc)
+
+        call_kwargs = mock_post.call_args.kwargs
+        data = call_kwargs["data"]
+        assert data["chunking_max_tokens"] == "512"
+        assert data["chunking_merge_peers"] == "false"
+        assert data["chunking_use_markdown_tables"] == "true"
+
+    @pytest.mark.asyncio
+    @patch("haiku.rag.chunkers.docling_serve.requests.post")
+    async def test_chunk_connection_error(self, mock_post, chunker):
+        """Test handling of connection errors."""
+        import requests
+
+        mock_post.side_effect = requests.exceptions.ConnectionError("Connection failed")
+
+        converter = get_converter(Config)
+        doc = converter.convert_text("# Test", name="test.md")
+
+        with pytest.raises(ValueError, match="Could not connect to docling-serve"):
+            await chunker.chunk(doc)
+
+    @pytest.mark.asyncio
+    @patch("haiku.rag.chunkers.docling_serve.requests.post")
+    async def test_chunk_timeout_error(self, mock_post, chunker):
+        """Test handling of timeout errors."""
+        import requests
+
+        mock_post.side_effect = requests.exceptions.Timeout("Timeout")
+
+        converter = get_converter(Config)
+        doc = converter.convert_text("# Test", name="test.md")
+
+        with pytest.raises(ValueError, match="timed out"):
+            await chunker.chunk(doc)
+
+    @pytest.mark.asyncio
+    @patch("haiku.rag.chunkers.docling_serve.requests.post")
+    async def test_chunk_auth_error(self, mock_post, chunker):
+        """Test handling of authentication errors."""
+        import requests
+
+        mock_response = Mock()
+        mock_response.status_code = 401
+        mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError(
+            response=mock_response
+        )
+        mock_post.return_value = mock_response
+
+        converter = get_converter(Config)
+        doc = converter.convert_text("# Test", name="test.md")
+
+        with pytest.raises(ValueError, match="Authentication failed"):
+            await chunker.chunk(doc)