From 81b7c7d9d125a22f0beba06dd1de945976f4ae44 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 17 Nov 2025 12:01:08 +0200 Subject: [PATCH] Tests for docling-serve chunker --- tests/test_chunker.py | 171 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 171 insertions(+) diff --git a/tests/test_chunker.py b/tests/test_chunker.py index f9bf5dde..205dd75c 100644 --- a/tests/test_chunker.py +++ b/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)