homelab-codex-ws/jobs/documents-ingest/tests/test_chunk_embed.py

"""Unit tests for the chunk + embed job — no DB, no real HTTP, no real Ollama.

`chunk_text`/`hard_split`/`split_paragraphs` tests moved to
packages/kb-mail/tests/test_chunking.py (module 5, faza mailowa, Krok 0) alongside the code."""
from __future__ import annotations

import json

import pytest

from documents_ingest.chunk_embed import (
    EmbeddingDimensionError,
    embed_chunk,
    extract_content,
    fetch_documents,
    fetch_existing_chunk_keys,
    insert_chunk,
    is_ocr_junk,
    run,
    _vector_literal,
)


def _paragraph(char: str, length: int) -> str:
    return char * length


class TestExtractContent:
    def test_finds_content_entity(self):
        entities = [{"type": "filename", "value": "a.pdf"}, {"type": "content", "text": "hello"}]
        assert extract_content(entities) == "hello"

    def test_missing_content_entity_returns_empty_string(self):
        entities = [{"type": "filename", "value": "a.pdf"}]
        assert extract_content(entities) == ""

    def test_none_text_returns_empty_string(self):
        entities = [{"type": "content", "text": None}]
        assert extract_content(entities) == ""

    def test_empty_entities_list(self):
        assert extract_content([]) == ""
        assert extract_content(None) == ""


class TestVectorLiteral:
    def test_formats_as_bracketed_csv(self):
        assert _vector_literal([0.1, 0.2, -0.3]) == "[0.1,0.2,-0.3]"


class TestIsOcrJunk:
    """Pilot cases from kb/phases/kb-m5-faza3.md §3.1 and the 2026-07-16 calibration
    session (kb/phases/kb-m5-eval-retrieval-pilot.md)."""

    def test_clean_text_is_not_junk(self):
        text = (
            "Ogólne Warunki Ubezpieczenia PZU Auto ustalają zasady odpowiedzialności "
            "ubezpieczyciela za szkody powstałe w związku z ruchem pojazdów mechanicznych."
        )
        assert is_ocr_junk(text) is False

    def test_ground_barcode_binary_noise_is_junk(self):
        # paperless:39/96/128/186/192 pattern: dense run of C0 control chars (scrambled
        # barcode), well above the calibrated raw-count threshold of 5.
        text = "\x01\x02\x03\x05\x06\x07\x0f\x10\x11\x12\x14\x15\x16\x17\x18\x19qQQQ\x11q\x01\x11!!AA"
        assert is_ocr_junk(text) is True

    def test_partial_mojibake_is_not_junk(self):
        # paperless:119: a confirmed retrieval hit (query 4, eval-set) despite partial
        # mojibake and a couple of stray control bytes — mojibake alone must not exclude it.
        text = (
            "Wygenerowano z systemu wydarzeń FIRST Israel | Turniej Regionalny FIRST LEGO "
            "League\n\x01integraln¹ czêœæ dokumentacji zawodów\nARTEFAKTY Z WĘDKI"
        )
        assert is_ocr_junk(text) is False

    def test_stray_single_control_char_in_legible_text_is_not_junk(self):
        # paperless:141 pattern: one stray control byte in an otherwise clean English
        # paragraph (a PDF-extraction artifact, not OCR noise).
        text = "separable\x0fconvolution is equal to the combination of a self-attention layer"
        assert is_ocr_junk(text) is False

    def test_dot_leader_toc_is_junk(self):
        text = ". . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ."
        assert is_ocr_junk(text) is True

    def test_empty_text_is_not_junk(self):
        assert is_ocr_junk("") is False


class _FakeEmbedResponse:
    def __init__(self, payload, status=200):
        self._payload = payload
        self._status = status

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False

    async def json(self):
        return self._payload

    def raise_for_status(self):
        if self._status >= 400:
            raise RuntimeError(f"HTTP {self._status}")


class _FakeOllamaSession:
    """Serves a fixed embedding vector for every /api/embeddings POST, or errors by filename."""

    def __init__(self, dim=1024, fail_for=None):
        self._dim = dim
        self._fail_for = fail_for or set()
        self.requests: list[dict] = []

    def post(self, url, json):
        self.requests.append({"url": url, "json": json})
        if json["prompt"] in self._fail_for:
            return _FakeEmbedResponse({}, status=500)
        return _FakeEmbedResponse({"embedding": [0.01] * self._dim})

    async def close(self):
        pass


class TestEmbedChunk:
    async def test_returns_embedding_and_elapsed(self):
        session = _FakeOllamaSession(dim=1024)
        embedding, elapsed = await embed_chunk(session, "http://fake-ollama", "bge-m3", "hello world")
        assert len(embedding) == 1024
        assert elapsed >= 0
        assert session.requests == [
            {"url": "http://fake-ollama/api/embeddings", "json": {"model": "bge-m3", "prompt": "hello world"}}
        ]

    async def test_missing_embedding_key_raises(self):
        class _EmptyResponse(_FakeEmbedResponse):
            pass

        class _Session:
            def post(self, url, json):
                return _EmptyResponse({})

        with pytest.raises(ValueError):
            await embed_chunk(_Session(), "http://fake-ollama", "bge-m3", "text")


class TestInsertSql:
    def test_on_conflict_target_includes_model(self):
        # Plan §3.3 fix: the old UNIQUE(envelope_id, chunk_index) silently no-opped a
        # second embedding model via ON CONFLICT DO NOTHING (review finding, phase 2 step
        # 6). The target must now include `model`.
        from documents_ingest.chunk_embed import _INSERT_SQL
        assert "ON CONFLICT (envelope_id, chunk_index, model)" in _INSERT_SQL


class TestInsertChunk:
    async def test_junk_chunk_inserts_null_embedding_with_reason(self):
        conn = _FakeConn()
        await insert_chunk(conn, "paperless:1", 0, "junk text", None, "bge-m3", excluded_reason="ocr_junk")

        assert conn.execute_calls == [("paperless:1", 0, "junk text", None, "bge-m3", "ocr_junk")]

    async def test_normal_chunk_inserts_vector_with_no_reason(self):
        conn = _FakeConn()
        await insert_chunk(conn, "paperless:1", 0, "real text", [0.1, 0.2], "bge-m3")

        assert conn.execute_calls == [("paperless:1", 0, "real text", "[0.1,0.2]", "bge-m3", None)]


class _FakeConn:
    def __init__(self, docs=None, existing_keys=None, existing_model="bge-m3", execute_results=None):
        self._docs = docs or []
        self._existing_keys = list(existing_keys or [])
        self._existing_model = existing_model
        # Optional queue of command tags returned by successive execute() calls, in order
        # (e.g. ["INSERT 0 0"] to simulate an ON CONFLICT DO NOTHING no-op). Defaults to a
        # real insert every time.
        self._execute_results = list(execute_results) if execute_results is not None else None
        self.execute_calls: list[tuple] = []

    async def fetch(self, query, *params):
        if "FROM document_chunk" in query:
            # Mirrors the real `WHERE model = $1` filter: existing_keys were "written" under
            # existing_model, so a query for a different model must not see them.
            if params and params[0] != self._existing_model:
                return []
            return [{"envelope_id": eid, "chunk_index": idx} for eid, idx in self._existing_keys]
        return self._docs

    async def execute(self, query, *params):
        self.execute_calls.append(params)
        if self._execute_results is not None:
            return self._execute_results.pop(0)
        return "INSERT 0 1"

    async def close(self):
        pass


def _row(envelope_id, content):
    return {"id": envelope_id, "entities": json.dumps([{"type": "content", "text": content}])}


class TestFetchHelpers:
    async def test_fetch_documents_no_limit_offset(self):
        conn = _FakeConn(docs=[_row("paperless:1", "hi")])
        docs = await fetch_documents(conn, limit=None, offset=None)
        assert docs == [_row("paperless:1", "hi")]

    async def test_fetch_existing_chunk_keys(self):
        conn = _FakeConn(existing_keys=[("paperless:1", 0), ("paperless:1", 1)])
        keys = await fetch_existing_chunk_keys(conn, model="bge-m3")
        assert keys == {("paperless:1", 0), ("paperless:1", 1)}


class TestRun:
    def _patch(self, monkeypatch, conn, ollama_session):
        async def _fake_connect(dsn):
            return conn
        monkeypatch.setattr("documents_ingest.chunk_embed.asyncpg.connect", _fake_connect)

        def _fake_session_factory(*args, **kwargs):
            return ollama_session
        monkeypatch.setattr("documents_ingest.chunk_embed.aiohttp.ClientSession", _fake_session_factory)

    async def test_dry_run_counts_without_calling_ollama_or_db(self, monkeypatch):
        conn = _FakeConn(docs=[_row("paperless:1", "short doc")])
        ollama = _FakeOllamaSession()
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=False)

        assert stats["documents_fetched"] == 1
        assert stats["documents_chunked"] == 1
        assert stats["chunks_total"] == 1
        assert stats["chunks_inserted"] == 1
        assert ollama.requests == []
        assert conn.execute_calls == []

    async def test_empty_content_counted_separately_not_as_error(self, monkeypatch):
        conn = _FakeConn(docs=[_row("paperless:1", ""), _row("paperless:2", "some text")])
        ollama = _FakeOllamaSession()
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=False)

        assert stats["empty_content"] == 1
        assert stats["documents_chunked"] == 1
        assert stats["documents_fetched"] == 2
        assert stats["chunks_errors"] == 0

    async def test_apply_embeds_and_inserts(self, monkeypatch):
        conn = _FakeConn(docs=[_row("paperless:1", "some real content")])
        ollama = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=True)

        assert stats["chunks_inserted"] == 1
        assert stats["embed_calls"] == 1
        assert len(conn.execute_calls) == 1
        assert len(ollama.requests) == 1

    async def test_idempotent_skips_already_embedded_chunks(self, monkeypatch):
        conn = _FakeConn(
            docs=[_row("paperless:1", "some real content")],
            existing_keys=[("paperless:1", 0)],
        )
        ollama = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=True)

        assert stats["chunks_already_embedded"] == 1
        assert stats["chunks_inserted"] == 0
        assert ollama.requests == []
        assert conn.execute_calls == []

    async def test_rerun_after_apply_inserts_nothing_new(self, monkeypatch):
        doc = _row("paperless:1", "some real content")

        conn1 = _FakeConn(docs=[doc])
        ollama1 = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn1, ollama1)
        first = await run(dsn="postgresql://fake", apply=True)
        assert first["chunks_inserted"] == 1

        conn2 = _FakeConn(docs=[doc], existing_keys=[("paperless:1", 0)])
        ollama2 = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn2, ollama2)
        second = await run(dsn="postgresql://fake", apply=True)

        assert second["chunks_inserted"] == 0
        assert second["chunks_already_embedded"] == 1
        assert ollama2.requests == []

    async def test_dimension_mismatch_aborts(self, monkeypatch):
        conn = _FakeConn(docs=[_row("paperless:1", "some real content")])
        ollama = _FakeOllamaSession(dim=768)  # wrong dim
        self._patch(monkeypatch, conn, ollama)

        with pytest.raises(EmbeddingDimensionError):
            await run(dsn="postgresql://fake", apply=True)

    async def test_embed_error_is_isolated_and_counted(self, monkeypatch):
        # Two chunks: force one to fail via a document long enough to produce 2 chunks
        # (chunk_size/overlap kept small to make the test fast and explicit). Real
        # word-like tokens (not a single 100-char run) so neither chunk trips is_ocr_junk.
        para1 = " ".join(["word"] * 20)
        para2 = " ".join(["term"] * 20)
        big = para1 + "\n\n" + para2
        conn = _FakeConn(docs=[_row("paperless:1", big)])
        ollama = _FakeOllamaSession(dim=1024, fail_for={para1})
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=True, chunk_size=100, chunk_overlap=20)

        assert stats["chunks_errors"] == 1
        assert stats["chunks_inserted"] == 1
        assert stats["chunks_total"] == (
            stats["chunks_already_embedded"] + stats["chunks_inserted"]
            + stats["chunks_conflict_skipped"] + stats["chunks_errors"]
        )

    async def test_insert_error_is_isolated_and_counted(self, monkeypatch):
        # Two chunks; the first chunk's INSERT raises (simulated DB blip) but the run
        # continues and the second chunk still gets embedded and inserted normally.
        # Real word-like tokens so neither chunk trips is_ocr_junk.
        big = " ".join(["word"] * 20) + "\n\n" + " ".join(["term"] * 20)
        conn = _FakeConn(docs=[_row("paperless:1", big)])

        real_execute = conn.execute

        async def _flaky_execute(query, *params):
            if params[1] == 0:  # chunk_index 0
                conn.execute_calls.append(params)
                raise RuntimeError("simulated db error")
            return await real_execute(query, *params)

        conn.execute = _flaky_execute

        ollama = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=True, chunk_size=100, chunk_overlap=20)

        assert stats["chunks_errors"] == 1
        assert stats["chunks_inserted"] == 1
        assert stats["chunks_total"] == (
            stats["chunks_already_embedded"] + stats["chunks_inserted"]
            + stats["chunks_conflict_skipped"] + stats["chunks_errors"]
        )

    async def test_conflict_skipped_counted_separately_from_inserted(self, monkeypatch):
        # ON CONFLICT DO NOTHING no-op: command tag reports 0 rows affected.
        conn = _FakeConn(
            docs=[_row("paperless:1", "some real content")], execute_results=["INSERT 0 0"]
        )
        ollama = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=True)

        assert stats["chunks_conflict_skipped"] == 1
        assert stats["chunks_inserted"] == 0
        assert stats["chunks_total"] == (
            stats["chunks_already_embedded"] + stats["chunks_inserted"]
            + stats["chunks_conflict_skipped"] + stats["chunks_errors"]
        )

    async def test_existing_keys_scoped_to_model(self, monkeypatch):
        # existing_keys were "written" under a different model — the WHERE model = $1
        # filter must not treat them as already-embedded for the current --model.
        conn = _FakeConn(
            docs=[_row("paperless:1", "some real content")],
            existing_keys=[("paperless:1", 0)],
            existing_model="some-other-model",
        )
        ollama = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=True, model="bge-m3")

        assert stats["chunks_already_embedded"] == 0
        assert stats["chunks_inserted"] == 1

    async def test_junk_chunk_flagged_without_calling_ollama(self, monkeypatch):
        junk = ". . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ."
        conn = _FakeConn(docs=[_row("paperless:1", junk)])
        ollama = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=True)

        assert stats["chunks_junk_flagged"] == 1
        assert stats["chunks_inserted"] == 0
        assert ollama.requests == []
        assert conn.execute_calls == [("paperless:1", 0, junk, None, "bge-m3", "ocr_junk")]
        assert stats["chunks_total"] == (
            stats["chunks_already_embedded"] + stats["chunks_inserted"] + stats["chunks_junk_flagged"]
            + stats["chunks_conflict_skipped"] + stats["chunks_errors"]
        )

    async def test_junk_chunk_dry_run_counted_without_db_or_ollama(self, monkeypatch):
        junk = ". . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ."
        conn = _FakeConn(docs=[_row("paperless:1", junk)])
        ollama = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=False)

        assert stats["chunks_junk_flagged"] == 1
        assert stats["chunks_inserted"] == 0
        assert ollama.requests == []
        assert conn.execute_calls == []

    async def test_junk_chunk_rerun_is_idempotent(self, monkeypatch):
        # Same (envelope_id, chunk_index) already flagged ocr_junk under this model — a
        # rerun must skip it, same as a normal embedded chunk (plan §3.1: backfill is a
        # one-shot, but chunk_embed.py itself must stay idempotent for new junk too).
        junk = ". . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ."
        conn = _FakeConn(
            docs=[_row("paperless:1", junk)],
            existing_keys=[("paperless:1", 0)],
        )
        ollama = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=True)

        assert stats["chunks_already_embedded"] == 1
        assert stats["chunks_junk_flagged"] == 0
        assert conn.execute_calls == []

    async def test_junk_conflict_skip_not_double_counted(self, monkeypatch):
        junk = ". . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ."
        conn = _FakeConn(docs=[_row("paperless:1", junk)], execute_results=["INSERT 0 0"])
        ollama = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=True)

        assert stats["chunks_conflict_skipped"] == 1
        assert stats["chunks_junk_flagged"] == 0
        assert ollama.requests == []

    async def test_stats_balance_documents_and_chunks(self, monkeypatch):
        conn = _FakeConn(docs=[
            _row("paperless:1", "content one"),
            _row("paperless:2", ""),
            _row("paperless:3", "content three"),
        ])
        ollama = _FakeOllamaSession(dim=1024)
        self._patch(monkeypatch, conn, ollama)

        stats = await run(dsn="postgresql://fake", apply=True)

        assert stats["documents_fetched"] == stats["empty_content"] + stats["documents_chunked"]
        assert stats["chunks_total"] == (
            stats["chunks_already_embedded"] + stats["chunks_inserted"]
            + stats["chunks_conflict_skipped"] + stats["chunks_errors"]
        )

    async def test_limit_and_offset_passed_through_to_query(self, monkeypatch):
        captured = {}

        class _Conn(_FakeConn):
            async def fetch(self, query, *params):
                if "FROM document_chunk" in query:
                    return []
                captured["query"] = query
                captured["params"] = params
                return []

        conn = _Conn()
        ollama = _FakeOllamaSession()
        self._patch(monkeypatch, conn, ollama)

        await run(dsn="postgresql://fake", apply=False, limit=10, offset=5)

        assert "LIMIT $1" in captured["query"]
        assert "OFFSET $2" in captured["query"]
        assert captured["params"] == (10, 5)