diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py index 57b2b23d..2e131401 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py @@ -143,7 +143,9 @@ class FSSource: if path.is_symlink(): try: resolved = path.resolve(strict=False) - except OSError: + except OSError: # pragma: no cover - strict=False absorbs + # symlink cycles and missing targets, so no real link + # reaches this; kept as a guard against platform drift. continue if not resolved.is_relative_to(self.root): continue diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 3871fd07..30c2ed11 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -244,6 +244,6 @@ def create_mcp_server( result = await rag.analyze(question, filter=filter, images=images) return result.answer except Exception as e: - return f"Error running analysis capability: {e!s}" # pragma: no cover + return f"Error running analysis capability: {e!s}" return mcp diff --git a/tests/ingester/test_fs_source.py b/tests/ingester/test_fs_source.py index 31c7b838..105b698b 100644 --- a/tests/ingester/test_fs_source.py +++ b/tests/ingester/test_fs_source.py @@ -330,3 +330,41 @@ async def test_fs_source_fetch_reads_off_event_loop_thread(fs_root: Path): "FSSource._read_body ran on the event-loop thread; the read+hash must " "be dispatched via asyncio.to_thread" ) + + +@pytest.mark.asyncio +async def test_fetch_rejects_foreign_scheme(tmp_path): + """`supports()` short-circuits on scheme, but fetch/head resolve directly, + so the unsupported-scheme path must be handled there too.""" + src = FSSource(root=tmp_path, supported_extensions=[".md"], source_id="local") + + with pytest.raises(UnsupportedSourceError): + await src.fetch("s3://bucket/key.md") + + assert await src.head("s3://bucket/key.md") is None + + +@pytest.mark.asyncio +async def test_fetch_falls_back_to_octet_stream_for_unknown_extension(tmp_path): + target = tmp_path / "data.unknownext" + target.write_bytes(b"payload") + src = FSSource( + root=tmp_path, supported_extensions=[".unknownext"], source_id="local" + ) + + result = await src.fetch(target.as_uri()) + + assert result.content_type == "application/octet-stream" + assert result.body == b"payload" + + +@pytest.mark.asyncio +async def test_discover_skips_symlink_to_missing_in_root_target(tmp_path): + """A broken symlink inside the root resolves to a path that is not a file.""" + (tmp_path / "real.md").write_text("real") + (tmp_path / "broken.md").symlink_to(tmp_path / "absent.md") + src = FSSource(root=tmp_path, supported_extensions=[".md"], source_id="local") + + events = [e async for e in src.discover()] + + assert {e.uri for e in events} == {(tmp_path / "real.md").as_uri()} diff --git a/tests/ingester/test_http_source.py b/tests/ingester/test_http_source.py index 08915e5e..856ad50c 100644 --- a/tests/ingester/test_http_source.py +++ b/tests/ingester/test_http_source.py @@ -422,3 +422,14 @@ async def test_fetch_skips_head_when_no_max_size(): ) await src.fetch("https://example.com/a.md") assert calls == ["GET"] + + +@pytest.mark.asyncio +async def test_aclose_closes_the_http_client(): + src = HTTPSource( + source_id="urls", + urls=[], + transport=httpx.MockTransport(lambda r: httpx.Response(200)), + ) + await src.aclose() + assert src._http.is_closed diff --git a/tests/ingester/test_pollers.py b/tests/ingester/test_pollers.py index d7970d13..b33bc2d3 100644 --- a/tests/ingester/test_pollers.py +++ b/tests/ingester/test_pollers.py @@ -602,3 +602,159 @@ async def test_fs_poller_enqueues_initial_files(tmp_path, jobs, sync): queued = await jobs.list_jobs(source_id="local") assert {Path(j.uri).name for j in queued} == {"a.md", "b.md"} assert all(j.status is JobStatus.QUEUED for j in queued) + + +# --- _dry_run_once --- + + +@pytest.mark.asyncio +async def test_dry_run_collects_changes_without_writing(fs_config, jobs, sync): + source = _StubSource( + "src", + [ + [ + _event("file:///a.md"), + _event("file:///b.md", kind=SourceEventKind.UNCHANGED), + _event("file:///c.md", kind=SourceEventKind.DELETE), + ] + ], + ) + poller = _periodic(source, fs_config, jobs, sync) + + ok, summary, changes = await poller._dry_run_once() + + assert ok is True + assert summary.upsert_count == 1 + assert summary.unchanged_count == 1 + assert summary.delete_count == 1 + assert {c.op for c in changes} == {JobOp.UPSERT, JobOp.DELETE} + # A dry run must not touch the queue. + assert await jobs.list_jobs(source_id="src") == [] + + +@pytest.mark.asyncio +async def test_dry_run_ignores_deletes_when_delete_orphans_false( + fs_config, jobs, sync, tmp_path +): + config = FSSourceConfig( + type="fs", + id="src", + root=tmp_path, + delete_orphans=False, + poll_interval_s=0.05, + ) + source = _StubSource("src", [[_event("file:///c.md", kind=SourceEventKind.DELETE)]]) + poller = _periodic(source, config, jobs, sync) + + ok, summary, changes = await poller._dry_run_once() + + assert ok is True + assert summary.delete_count == 0 + assert summary.ignored_delete_count == 1 + assert changes == [] + + +@pytest.mark.asyncio +async def test_dry_run_skipped_when_circuit_open(fs_config, jobs, sync): + class _Clock: + now = 0.0 + + def __call__(self): + return self.now + + breaker = CircuitBreaker( + CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0), + now_fn=_Clock(), + ) + source = _StubSource("src", []) + source.fail_with = RuntimeError("upstream down") + poller = _periodic(source, fs_config, jobs, sync, breaker=breaker) + + assert await poller._sweep_once() is False + assert breaker.is_open is True + + before = source.discover_calls + ok, summary, changes = await poller._dry_run_once() + + assert ok is False + assert changes == [] + assert source.discover_calls == before + assert poller.last_skip_reason == "circuit_open" + + +@pytest.mark.asyncio +async def test_dry_run_records_failure_when_discover_raises(fs_config, jobs, sync): + source = _StubSource("src", []) + source.fail_with = RuntimeError("upstream down") + poller = _periodic(source, fs_config, jobs, sync) + + ok, summary, changes = await poller._dry_run_once() + + assert ok is False + assert changes == [] + assert poller._breaker.consecutive_failures == 1 + + +@pytest.mark.asyncio +async def test_dry_run_skipped_when_queue_has_pending_work(fs_config, jobs, sync): + source = _StubSource("src", [[_event("file:///a.md")]]) + poller = _periodic(source, fs_config, jobs, sync) + await jobs.enqueue("src", "file:///pending.md", JobOp.UPSERT) + + ok, _summary, changes = await poller._dry_run_once() + + assert ok is False + assert changes == [] + assert poller.last_skip_reason == "pending_work" + + +@pytest.mark.asyncio +async def test_watch_deleted_skipped_when_delete_orphans_false(tmp_path, jobs, sync): + from watchfiles import Change + + from haiku.rag.ingester.pollers.fs import FSPoller + from haiku.rag.ingester.sources.fs import FSSource + + cfg = FSSourceConfig( + type="fs", + id="local", + root=tmp_path, + delete_orphans=False, + poll_interval_s=60.0, + ) + poller = FSPoller( + source=FSSource(root=tmp_path, supported_extensions=[".md"], source_id="local"), + config=cfg, + job_repo=jobs, + sync_repo=sync, + ) + + await poller._handle_watch_change(Change.deleted, tmp_path / "gone.md") + + assert await jobs.list_jobs(source_id="local") == [] + + +@pytest.mark.asyncio +async def test_dry_run_manifest_reports_failed_sources(tmp_path, jobs, sync): + """A source whose discover() raises is named in the failed list while the + manifest still carries the sources that succeeded.""" + manager = PollerManager( + configs=[FSSourceConfig(type="fs", id="ok", root=tmp_path)], + job_repo=jobs, + sync_repo=sync, + ) + broken = _StubSource("broken", []) + broken.fail_with = RuntimeError("upstream down") + manager._pollers.append( + _periodic( + broken, + FSSourceConfig(type="fs", id="broken", root=tmp_path, poll_interval_s=60.0), + jobs, + sync, + ) + ) + + manifest, failed = await manager.dry_run_manifest() + + assert failed == ["broken"] + assert {s.source_id for s in manifest.sources} == {"ok", "broken"} diff --git a/tests/ingester/test_webdav_source.py b/tests/ingester/test_webdav_source.py index 060ddb06..1da7e450 100644 --- a/tests/ingester/test_webdav_source.py +++ b/tests/ingester/test_webdav_source.py @@ -7,26 +7,29 @@ from haiku.rag.ingester.sources.base import FileTooLargeError, SourceEventKind from haiku.rag.ingester.sources.webdav import WebDAVSource, _strip_etag -def test_strip_etag_strong_quoted(): - assert _strip_etag('"abc123"') == "abc123" - - -def test_strip_etag_weak_marker(): - assert _strip_etag('W/"abc123"') == "abc123" - - -def test_strip_etag_unquoted(): - assert _strip_etag("abc123") == "abc123" - - -def test_strip_etag_whitespace(): - assert _strip_etag(' W/"abc" ') == "abc" - - -def test_strip_etag_empty_returns_none(): - assert _strip_etag("") is None - assert _strip_etag('""') is None - assert _strip_etag(None) is None +@pytest.mark.parametrize( + "raw,expected", + [ + ('"abc123"', "abc123"), + ('W/"abc123"', "abc123"), + ("abc123", "abc123"), + (' W/"abc" ', "abc"), + ("", None), + ('""', None), + (None, None), + ], + ids=[ + "strong_quoted", + "weak_marker", + "unquoted", + "whitespace", + "empty", + "empty_quotes", + "none", + ], +) +def test_strip_etag(raw, expected): + assert _strip_etag(raw) == expected def _transport(handler) -> httpx.MockTransport: @@ -628,3 +631,131 @@ async def test_fetch_skips_head_when_no_max_size(): ) await src.fetch("https://nc.example.com/dav/a.txt") assert calls == ["GET"] + + +# Malformed multistatus bodies: a that can't be decoded is dropped +# rather than aborting the whole listing. + + +def _raw_multistatus(*response_blocks: str) -> bytes: + body = ['', ''] + body.extend(response_blocks) + body.append("") + return "\n".join(body).encode() + + +_NO_HREF = """ + + HTTP/1.1 200 OK + "r" + + """ + +_EMPTY_HREF = """ + + + HTTP/1.1 200 OK + "r" + + """ + +_NO_STATUS = """ + /dav/a.md + + "r" + + """ + +_NOT_FOUND_STATUS = """ + /dav/a.md + + HTTP/1.1 404 Not Found + "r" + + """ + +_STATUS_WITHOUT_PROP = """ + /dav/a.md + + HTTP/1.1 200 OK + + """ + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "block", + [_NO_HREF, _EMPTY_HREF, _NO_STATUS, _NOT_FOUND_STATUS], + ids=["no_href", "empty_href", "propstat_without_status", "propstat_404"], +) +async def test_head_returns_none_for_undecodable_response(block): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(207, content=_raw_multistatus(block)) + + src = WebDAVSource( + source_id="nc", + base_url="https://nc.example.com/dav/", + transport=_transport(handler), + ) + assert await src.head("https://nc.example.com/dav/a.md") is None + + +@pytest.mark.asyncio +async def test_head_returns_none_for_empty_multistatus(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(207, content=_raw_multistatus()) + + src = WebDAVSource( + source_id="nc", + base_url="https://nc.example.com/dav/", + transport=_transport(handler), + ) + assert await src.head("https://nc.example.com/dav/a.md") is None + + +@pytest.mark.asyncio +async def test_entry_with_status_but_no_prop_has_no_revision(): + """A 200 propstat carrying no still yields an entry, without a revision.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(207, content=_raw_multistatus(_STATUS_WITHOUT_PROP)) + + src = WebDAVSource( + source_id="nc", + base_url="https://nc.example.com/dav/", + transport=_transport(handler), + ) + assert await src.head("https://nc.example.com/dav/a.md") is None + + +@pytest.mark.asyncio +async def test_discover_skips_base_url_reported_as_file(): + """Broken servers list the base URL itself as a non-collection; it and any + href outside the base are skipped.""" + body = _multistatus( + {"href": "/dav/", "etag": '"base"'}, + {"href": "/outside/x.md", "etag": '"out"'}, + {"href": "/dav/keep.md", "etag": '"keep"'}, + ) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(207, content=body) + + src = WebDAVSource( + source_id="nc", + base_url="https://nc.example.com/dav/", + transport=_transport(handler), + ) + events = [event async for event in src.discover()] + assert {e.uri for e in events} == {"https://nc.example.com/dav/keep.md"} + + +@pytest.mark.asyncio +async def test_aclose_closes_the_http_client(): + src = WebDAVSource( + source_id="nc", + base_url="https://nc.example.com/dav/", + transport=_transport(lambda r: httpx.Response(200)), + ) + await src.aclose() + assert src._http.is_closed diff --git a/tests/test_chunker.py b/tests/test_chunker.py index b795bdea..97b0514c 100644 --- a/tests/test_chunker.py +++ b/tests/test_chunker.py @@ -798,3 +798,78 @@ async def test_serve_chunker_accepts_picture_laden_docling(doclaynet_first_page_ chunks = await serve_chunker.chunk(doc) assert len(chunks) > 0, "docling-serve chunker returned 0 chunks" + + +class TestDoclingServeChunkerRefResolution: + """_resolve_label_from_document and the dict-shaped doc_items branch.""" + + @pytest.fixture + def chunker(self): + config = AppConfig() + config.providers.docling_serve.base_url = "http://localhost:5001" + config.processing.chunk_size = 256 + config.processing.chunking_tokenizer = "Qwen/Qwen3-Embedding-0.6B" + return DoclingServeChunker(config) + + @pytest.fixture + def document(self): + from docling_core.types.doc.document import DoclingDocument + + return DoclingDocument.model_validate( + { + "name": "doc", + "texts": [ + { + "self_ref": "#/texts/0", + "text": "body", + "orig": "body", + "label": "paragraph", + } + ], + "tables": [], + "pictures": [], + "groups": [], + "body": {"self_ref": "#/body", "children": []}, + "furniture": {"self_ref": "#/furniture", "children": []}, + } + ) + + @pytest.mark.parametrize( + "ref", + ["not-a-ref", "#/texts/999", "#/nope/0"], + ids=["unparseable", "index_out_of_range", "unknown_collection"], + ) + def test_unresolvable_ref_yields_no_label(self, document, ref): + from haiku.rag.chunkers.docling_serve import _resolve_label_from_document + + assert _resolve_label_from_document(ref, document) is None + + def test_resolvable_ref_yields_label(self, document): + from haiku.rag.chunkers.docling_serve import _resolve_label_from_document + + assert _resolve_label_from_document("#/texts/0", document) == "paragraph" + + @pytest.mark.asyncio + async def test_chunk_of_none_returns_empty(self, chunker): + assert await chunker.chunk(None) == [] + + @pytest.mark.asyncio + async def test_dict_shaped_doc_items_are_decoded(self, chunker, document): + """docling-serve returns refs as strings today; the dict shape is + accepted in case the API changes.""" + + async def fake_chunk_api(_document): + return [ + { + "raw_text": "body", + "doc_items": [{"self_ref": "#/texts/0", "label": "paragraph"}], + } + ] + + chunker._call_chunk_api = fake_chunk_api # type: ignore[method-assign] + + chunks = await chunker.chunk(document) + + assert len(chunks) == 1 + assert chunks[0].metadata["doc_item_refs"] == ["#/texts/0"] + assert chunks[0].metadata["labels"] == ["paragraph"] diff --git a/tests/test_converters.py b/tests/test_converters.py index 74518b12..9d686227 100644 --- a/tests/test_converters.py +++ b/tests/test_converters.py @@ -405,6 +405,46 @@ class TestDoclingLocalConverter: assert isinstance(doc, DoclingDocument) assert doc.name == "test" + @pytest.mark.asyncio + async def test_convert_file_reads_unknown_extension_as_text( + self, converter, tmp_path + ): + """An extension in neither the docling nor the text set is read as text.""" + source = tmp_path / "notes.xyz" + source.write_text("Plain body for an unknown extension.") + + doc = await converter.convert_file(source) + + assert isinstance(doc, DoclingDocument) + assert "Plain body for an unknown extension." in doc.export_to_markdown() + + @pytest.mark.asyncio + async def test_convert_file_raises_for_undecodable_file(self, converter, tmp_path): + source = tmp_path / "binary.xyz" + source.write_bytes(b"\xff\xfe\x00\x01 not utf-8") + + with pytest.raises(ValueError, match="Failed to parse file"): + await converter.convert_file(source) + + @pytest.mark.asyncio + async def test_convert_text_wraps_conversion_failure(self, converter, monkeypatch): + def boom(*_args, **_kwargs): + raise RuntimeError("docling exploded") + + monkeypatch.setattr(converter, "_sync_convert_docling_text", boom) + + with pytest.raises(ValueError, match="Failed to convert text"): + await converter.convert_text("# Test", name="test.md") + + @pytest.mark.asyncio + async def test_convert_text_falls_back_when_format_not_inferable(self, converter): + """docling raises ConversionError for an extension it has no backend + for; the simple-document fallback keeps the text.""" + doc = await converter.convert_text("just some prose", name="mystery.zzz") + + assert isinstance(doc, DoclingDocument) + assert "just some prose" in doc.export_to_markdown() + @pytest.mark.asyncio async def test_convert_code_file(self, converter): """Test that code files are wrapped in code blocks.""" @@ -1522,3 +1562,121 @@ class TestDoclingServeConverterIntegration: assert str(sample.image.uri).startswith("data:image/"), ( "Rehydrated picture URI should be a data: URI, not a bare artifact filename" ) + + +class TestDoclingServeZipParsing: + """_parse_zip_to_docling decodes the target_type=zip payload. These drive + its branches directly — no docling-serve instance involved.""" + + @pytest.fixture + def converter(self): + config = AppConfig() + config.processing.converter = "docling-serve" + conv = get_converter(config) + assert isinstance(conv, DoclingServeConverter) + return conv + + @staticmethod + def _zip(entries: dict[str, bytes]) -> bytes: + import io + import zipfile + + buf = io.BytesIO() + with zipfile.ZipFile(buf, mode="w") as zf: + for name, blob in entries.items(): + zf.writestr(name, blob) + return buf.getvalue() + + @staticmethod + def _doc_json(**extra) -> dict: + base = { + "name": "document", + "texts": [], + "tables": [], + "pictures": [], + "groups": [], + "body": {"self_ref": "#/body", "children": []}, + "furniture": {"self_ref": "#/furniture", "children": []}, + } + base.update(extra) + return base + + def test_raises_without_top_level_json(self, converter): + blob = self._zip({"artifacts/image.png": b"png"}) + + with pytest.raises(ValueError, match="no top-level JSON document"): + converter._parse_zip_to_docling(blob, "doc.pdf") + + def test_picture_without_image_is_left_alone(self, converter): + import json as _json + + doc_json = self._doc_json( + pictures=[ + { + "self_ref": "#/pictures/0", + "label": "picture", + "image": None, + "prov": [], + } + ] + ) + blob = self._zip({"document.json": _json.dumps(doc_json).encode()}) + + doc = converter._parse_zip_to_docling(blob, "doc.pdf") + + assert doc.pictures[0].image is None + + def test_data_uri_image_is_passed_through(self, converter): + import json as _json + + data_uri = "data:image/png;base64,aGVsbG8=" + doc_json = self._doc_json( + pictures=[ + { + "self_ref": "#/pictures/0", + "label": "picture", + "image": { + "mimetype": "image/png", + "dpi": 72, + "size": {"width": 1, "height": 1}, + "uri": data_uri, + }, + "prov": [], + } + ] + ) + blob = self._zip({"document.json": _json.dumps(doc_json).encode()}) + + doc = converter._parse_zip_to_docling(blob, "doc.pdf") + + assert str(doc.pictures[0].image.uri) == data_uri + + def test_non_dict_page_entry_is_skipped_while_inlining(self, converter): + """A page entry that isn't an object must not blow up the image-inlining + loop with an AttributeError; it falls through to schema validation.""" + import json as _json + + from pydantic import ValidationError + + doc_json = self._doc_json(pages={"1": "not-a-page-object"}) + blob = self._zip({"document.json": _json.dumps(doc_json).encode()}) + + with pytest.raises(ValidationError): + converter._parse_zip_to_docling(blob, "doc.pdf") + + @pytest.mark.asyncio + async def test_convert_text_rejects_unsupported_format(self, converter): + with pytest.raises(ValueError, match="Unsupported format"): + await converter.convert_text("body", format="pdf") + + @pytest.mark.asyncio + async def test_convert_text_plain_builds_document_locally(self, converter): + """format="plain" never reaches the network.""" + converter.client.submit_and_poll_zip = AsyncMock( # ty: ignore[invalid-assignment] + side_effect=AssertionError("must not call docling-serve") + ) + + doc = await converter.convert_text("just text", format="plain") + + assert isinstance(doc, DoclingDocument) + assert "just text" in doc.export_to_markdown() diff --git a/tests/test_docling_serve_client.py b/tests/test_docling_serve_client.py index ceaf89ee..175ed53e 100644 --- a/tests/test_docling_serve_client.py +++ b/tests/test_docling_serve_client.py @@ -481,3 +481,29 @@ def test_from_config_wires_retry_and_breaker(): assert client._max_attempts == 7 assert client._breaker_config.failure_threshold == 9 assert client._breaker_config.cooldown_s == 90.0 + + +@pytest.mark.asyncio +async def test_submit_without_task_id_raises(): + """A 200 that carries no task_id is a protocol violation, not a silent pass.""" + import httpx + + from haiku.rag.providers.docling_serve import DoclingServeClient + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={}) + + client = DoclingServeClient(base_urls="http://docling:5001") + files = {"files": ("doc.pdf", b"pdf", "application/octet-stream")} + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http: + with pytest.raises(ValueError, match="did not return a task_id"): + await client._submit_and_wait( + http, + "http://docling:5001", + "/v1/convert/source/async", + files, + {}, + {}, + "doc.pdf", + ) diff --git a/tests/test_embedder_config.py b/tests/test_embedder_config.py index a8a37fde..ce7561ff 100644 --- a/tests/test_embedder_config.py +++ b/tests/test_embedder_config.py @@ -148,3 +148,47 @@ def test_vllm_embedder_does_not_double_append_v1(): base_url = embedder._base_url.rstrip("/") # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] assert base_url.endswith("/v1") assert not base_url.endswith("/v1/v1") + + +def test_vector_dim_property_reports_configured_dimension(): + from haiku.rag.embeddings import EmbedderWrapper + + assert EmbedderWrapper(embedder=None, vector_dim=512).vector_dim == 512 + + +@pytest.mark.parametrize( + "provider,env_var", + [("voyageai", "VOYAGE_API_KEY"), ("cohere", "CO_API_KEY")], +) +def test_saas_providers_build_offline(monkeypatch, provider, env_var): + """Construction only wires the SDK; no request is made.""" + monkeypatch.setenv(env_var, "test-key") + config = AppConfig( + embeddings=EmbeddingsConfig( + model=EmbeddingModelConfig( + provider=provider, name="some-model", vector_dim=1024 + ), + ), + ) + + embedder = get_embedder(config) + + assert embedder.vector_dim == 1024 + + +def test_cohere_floats_rejects_missing_embeddings(): + from types import SimpleNamespace + + from haiku.rag.embeddings.cohere import _floats + + result = SimpleNamespace(embeddings=SimpleNamespace(float_=None)) + + with pytest.raises(ValueError, match="no float embeddings"): + _floats(result) + + +def test_voyageai_to_pil_rejects_unsupported_type(): + from haiku.rag.embeddings.voyageai import _to_pil + + with pytest.raises(TypeError, match="Unsupported image type"): + _to_pil("not an image") # ty: ignore[invalid-argument-type] diff --git a/tests/test_mcp.py b/tests/test_mcp.py index eda2e2e2..f0c36474 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -332,3 +332,152 @@ class TestMCPImageInput: result = await ask(question="q") assert result == "answer" assert captured["images"] is None + + +class TestMCPFileAndUrlIngestion: + @pytest.mark.asyncio + async def test_add_document_from_file(self, temp_db_path, tmp_path): + async with HaikuRAG(temp_db_path, create=True): + pass + source = tmp_path / "note.txt" + source.write_text("Ingested from a file path.") + + mcp = create_mcp_server(temp_db_path, read_only=False) + add_file = await _get_tool(mcp, "add_document_from_file") + + doc_id = await add_file(file_path=str(source), title="File Doc") + assert doc_id is not None + + get_doc = await _get_tool(mcp, "get_document") + doc = await get_doc(document_id=doc_id) + assert doc.title == "File Doc" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "tool_name,kwargs", + [ + ("add_document_from_file", {"file_path": "/tmp/x.txt"}), + ("add_document_from_url", {"url": "https://example.com/x.txt"}), + ], + ) + @pytest.mark.parametrize( + "results,expected", + [ + ( + [Document(id="first", content="a"), Document(id="second", content="b")], + "first", + ), + ([], None), + ], + ids=["directory_reports_first_id", "empty_directory_reports_none"], + ) + async def test_add_tools_handle_multi_document_sources( + self, mcp_db, monkeypatch, tool_name, kwargs, results, expected + ): + """A source resolving to several documents reports the first id.""" + + async def fake_from_source(self, source, title=None, metadata=None, **kw): + return results + + monkeypatch.setattr(HaikuRAG, "create_document_from_source", fake_from_source) + mcp = create_mcp_server(mcp_db, read_only=False) + add = await _get_tool(mcp, tool_name) + + assert await add(**kwargs) == expected + + @pytest.mark.asyncio + async def test_add_document_from_url(self, mcp_db, monkeypatch): + async def fake_from_source(self, source, title=None, metadata=None, **kwargs): + assert source == "https://example.com/doc.txt" + return Document(id="url-doc", content="fetched") + + monkeypatch.setattr(HaikuRAG, "create_document_from_source", fake_from_source) + mcp = create_mcp_server(mcp_db, read_only=False) + add_url = await _get_tool(mcp, "add_document_from_url") + + assert await add_url(url="https://example.com/doc.txt") == "url-doc" + + +class TestMCPToolsDegradeOnError: + """Every tool swallows client failures and returns its empty value rather + than propagating an exception to the MCP transport.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "client_method,tool_name,kwargs,expected", + [ + ( + "create_document_from_source", + "add_document_from_file", + {"file_path": "/tmp/x.txt"}, + None, + ), + ( + "create_document_from_source", + "add_document_from_url", + {"url": "https://example.com/x"}, + None, + ), + ("create_document", "add_document_from_text", {"content": "x"}, None), + ("delete_document", "delete_document", {"document_id": "x"}, False), + ("search", "search_documents", {"query": "x"}, []), + ("get_document_by_id", "get_document", {"document_id": "x"}, None), + ("list_documents", "list_documents", {}, []), + ], + ) + async def test_tool_returns_empty_value_when_client_raises( + self, mcp_db, monkeypatch, client_method, tool_name, kwargs, expected + ): + async def boom(self, *args, **kw): + raise RuntimeError("client exploded") + + monkeypatch.setattr(HaikuRAG, client_method, boom) + mcp = create_mcp_server(mcp_db, read_only=False) + tool = await _get_tool(mcp, tool_name) + + assert await tool(**kwargs) == expected + + @pytest.mark.asyncio + async def test_list_documents_returns_empty_for_invalid_filter(self, mcp_db): + mcp = create_mcp_server(mcp_db, read_only=True) + list_docs = await _get_tool(mcp, "list_documents") + + assert await list_docs(filter="no_such_column = 1") == [] + + @pytest.mark.asyncio + async def test_analyze_reports_the_error(self, mcp_db, monkeypatch): + async def boom(self, question, filter=None, images=None): + raise RuntimeError("sandbox exploded") + + monkeypatch.setattr(HaikuRAG, "analyze", boom) + mcp = create_mcp_server(mcp_db, read_only=True) + analyze = await _get_tool(mcp, "analyze") + + assert "sandbox exploded" in await analyze(question="q") + + @pytest.mark.asyncio + async def test_ask_question_appends_citations_when_requested( + self, mcp_db, monkeypatch + ): + from haiku.rag.store.models.citation import Citation + + citation = Citation( + chunk_id="c1", + document_id="d1", + content="cited text", + document_uri="test://ai-overview", + document_title="AI Overview", + ) + + async def fake_ask(self, question, filter=None, images=None): + return ("the answer", [citation]) + + monkeypatch.setattr(HaikuRAG, "ask", fake_ask) + mcp = create_mcp_server(mcp_db, read_only=True) + ask = await _get_tool(mcp, "ask_question") + + with_cite = await ask(question="q", cite=True) + assert with_cite.startswith("the answer") + assert "AI Overview" in with_cite + + assert await ask(question="q", cite=False) == "the answer" diff --git a/tests/test_pdf_split.py b/tests/test_pdf_split.py index 09b3f274..af5abaf9 100644 --- a/tests/test_pdf_split.py +++ b/tests/test_pdf_split.py @@ -243,3 +243,14 @@ def test_concatenate_shifts_page_nos_and_unique_self_refs(): assert page_nos == [1, 2], f"expected b's page 1 to shift to page 2, got {page_nos}" assert sorted(merged.pages.keys()) == [1, 2] + + +def test_iter_pdf_slices_rejects_unopenable_pdf(tmp_path): + """pdfium refuses non-PDF bytes; the caller sees UnsupportedSourceError.""" + from haiku.rag.client.exceptions import UnsupportedSourceError + + junk = tmp_path / "not-really.pdf" + junk.write_bytes(b"this is not a pdf at all") + + with pytest.raises(UnsupportedSourceError, match="cannot open PDF"): + list(iter_pdf_slices(junk, slice_size=1)) diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 0f8c77bf..635fe5fc 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -479,3 +479,30 @@ async def test_cross_encoder_reranker(): assert "0" in top_ids or "2" in top_ids except ImportError: pytest.skip("sentence-transformers not installed") + + +@pytest.mark.asyncio +async def test_cross_encoder_reranks_via_model_ranking(monkeypatch): + """The rank() results map back onto the input chunks by corpus_id.""" + from haiku.rag.reranking import cross_encoder as ce_module + + class _StubCrossEncoder: + def __init__(self, model): + self.model = model + + def rank(self, query, documents, top_k=10): + # Reverse order so the mapping back to chunks is observable. + return [ + {"corpus_id": i, "score": 1.0 - (i / 10)} + for i in reversed(range(len(documents))) + ][:top_k] + + monkeypatch.setattr(ce_module, "CrossEncoder", _StubCrossEncoder) + + reranker = ce_module.CrossEncoderReranker("stub/model") + reranked = await reranker.rerank("query", chunks, top_n=2) + + assert len(reranked) == 2 + last_index = len(chunks) - 1 + assert reranked[0][0] is chunks[last_index] + assert reranked[0][1] == pytest.approx(1.0 - last_index / 10)