DocsGPT/tests/scripts/test_reembed.py

"""Re-embed script: CLI contract, orchestration, and the FAISS rebuild path."""

from unittest.mock import MagicMock, patch

import pytest

from docsgpt.scripts import reembed


def paginating_cursor(chunk_rows, *, graph_rows=(), graph_table=("graph_nodes",)):
    """A cursor answering the reads ``reembed_pgvector`` issues, from memory.

    The chunk read is paginated by keyset, so a cursor that returns the same
    page for every ``fetchall`` never terminates. Returns the cursor and the
    mutable table backing it.

    Args:
        chunk_rows: ``(id, text)`` rows for the source's chunk table.
        graph_rows: ``(id, name)`` rows for ``graph_nodes``.
        graph_table: What ``to_regclass`` reports; ``(None,)`` for absent.

    Returns:
        ``(cursor, tables)``, where ``tables["chunks"]`` and
        ``tables["graph"]`` can be reassigned to change what is served.
    """
    tables = {"chunks": list(chunk_rows), "graph": list(graph_rows)}
    pending = {"rows": []}
    cursor = MagicMock()

    def execute(query, params=None):
        text = str(query)
        if "to_regclass" in text:
            pending["rows"] = [graph_table]
        elif "FROM graph_nodes" in text:
            pending["rows"] = list(tables["graph"])
        elif "count(*)" in text:
            pending["rows"] = [(len(tables["chunks"]),)]
        elif "id > %s" in text:
            _, after_id, limit = params
            pending["rows"] = [r for r in tables["chunks"] if r[0] > after_id][:limit]
        else:
            pending["rows"] = list(tables["chunks"])[: params[1]]

    cursor.execute.side_effect = execute
    cursor.fetchall.side_effect = lambda: pending["rows"]
    cursor.fetchone.side_effect = lambda: pending["rows"][0] if pending["rows"] else None
    return cursor, tables


class TestCLI:
    def test_defaults(self):
        args = reembed.build_parser().parse_args([])
        assert args.dry_run is False
        assert args.sources is None
        assert args.batch_size == reembed.DEFAULT_BATCH_SIZE

    def test_unsupported_store_exits_without_touching_anything(self):
        with patch.object(reembed.settings, "VECTOR_STORE", "qdrant", create=True):
            with patch.object(reembed, "run") as run:
                assert reembed.main([]) == 2
        run.assert_not_called()

    def test_supported_stores_are_pgvector_and_faiss(self):
        assert set(reembed.SUPPORTED_STORES) == {"pgvector", "faiss"}

    def test_sources_are_split_and_trimmed(self):
        with patch.object(reembed.settings, "VECTOR_STORE", "faiss", create=True):
            with patch.object(reembed, "run", return_value=0) as run:
                reembed.main(["--sources", " a , b ,, c "])
        assert run.call_args.args[1] == ["a", "b", "c"]

    def test_batch_size_is_clamped_to_at_least_one(self):
        with patch.object(reembed.settings, "VECTOR_STORE", "faiss", create=True):
            with patch.object(reembed, "run", return_value=0) as run:
                reembed.main(["--batch-size", "0"])
        assert run.call_args.args[2] == 1


class TestRun:
    def test_no_sources_is_a_clean_exit(self):
        with patch.object(reembed, "list_source_ids", return_value=[]):
            assert reembed.run("faiss", None, 8, False) == 0

    def test_processes_every_discovered_source(self):
        with patch.object(reembed, "list_source_ids", return_value=["a", "b"]):
            with patch.object(reembed, "reembed_faiss", return_value=(3, 3)) as handler:
                assert reembed.run("faiss", None, 8, False) == 0
        assert [call.args[0] for call in handler.call_args_list] == ["a", "b"]

    def test_explicit_sources_skip_discovery(self):
        with patch.object(reembed, "list_source_ids") as discover:
            with patch.object(reembed, "reembed_faiss", return_value=(1, 1)):
                reembed.run("faiss", ["only-this"], 8, False)
        discover.assert_not_called()

    def test_one_failing_source_does_not_stop_the_others(self):
        def handler(source_id, batch_size, dry_run):
            if source_id == "bad":
                raise RuntimeError("boom")
            return (1, 1)

        with patch.object(reembed, "reembed_faiss", side_effect=handler) as spy:
            code = reembed.run("faiss", ["good", "bad", "also-good"], 8, False)
        assert code == 1, "a failure must be reported in the exit code"
        assert spy.call_count == 3, "later sources must still be attempted"

    def test_dry_run_is_reported_as_success(self):
        with patch.object(reembed, "reembed_faiss", return_value=(5, 0)):
            assert reembed.run("faiss", ["a"], 8, True) == 0

    def test_pgvector_uses_the_pgvector_handler(self):
        with patch.object(reembed, "reembed_pgvector", return_value=(1, 1)) as handler:
            reembed.run("pgvector", ["a"], 8, False)
        handler.assert_called_once()


class TestFaissRebuild:
    @pytest.fixture
    def stores(self):
        """An existing store to read from and the rebuilt one written back."""
        existing = MagicMock()
        existing.get_chunks.return_value = [
            {"doc_id": "1", "text": "alpha", "metadata": {"i": 0}},
            {"doc_id": "2", "text": "beta", "metadata": {"i": 1}},
        ]
        rebuilt = MagicMock()
        with patch.object(
            reembed.VectorCreator, "create_vectorstore", side_effect=[existing, rebuilt]
        ) as factory:
            yield existing, rebuilt, factory

    def test_rebuilds_from_stored_text_and_saves(self, stores):
        existing, rebuilt, factory = stores
        seen, written = reembed.reembed_faiss("s1", batch_size=8, dry_run=False)

        assert (seen, written) == (2, 2)
        rebuilt.save_local.assert_called_once()
        docs = factory.call_args_list[1].kwargs["docs_init"]
        assert [d.page_content for d in docs] == ["alpha", "beta"]
        assert [d.metadata for d in docs] == [{"i": 0}, {"i": 1}]

    def test_embeddings_key_comes_from_settings(self, stores):
        """A placeholder here is sent as the server's bearer token."""
        _, _, factory = stores
        with patch.object(reembed.settings, "EMBEDDINGS_KEY", "sk-real", create=True):
            reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
        keys = [call.kwargs.get("embeddings_key") for call in factory.call_args_list]
        assert keys == ["sk-real", "sk-real"]

    def test_chunk_ids_are_preserved(self, stores):
        """Re-embedding must not renumber chunks.

        Fresh ids orphan every GraphRAG ``graph_node_chunks`` row for the
        source and invalidate any id a client already holds.
        """
        _, _, factory = stores
        reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
        assert factory.call_args_list[1].kwargs["ids"] == ["1", "2"]

    def test_existing_index_is_opened_past_the_dimension_check(self, stores):
        """A width change is the main reason to run this script.

        ``assert_embedding_dimensions`` refuses to open an index whose width
        differs from the configured model -- and its error message recommends
        this script, so without the opt-out the advice failed on every source.
        The chunk text lives in the sidecar, so reading it needs no match.
        """
        _, _, factory = stores
        reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
        assert factory.call_args_list[0].kwargs["skip_dimension_check"] is True
        # The rebuild writes the new width, so it must still be checked.
        assert "skip_dimension_check" not in factory.call_args_list[1].kwargs

    def test_batch_size_is_forwarded_to_the_rebuild(self, stores):
        """On a remote embeddings server the whole index is otherwise one POST."""
        _, _, factory = stores
        reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
        assert factory.call_args_list[1].kwargs["batch_size"] == 8

    def test_existing_index_is_not_deleted(self, stores):
        """The rebuild must not destroy the old index before the new one exists."""
        existing, rebuilt, _ = stores
        reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
        existing.delete_index.assert_not_called()

    def test_dry_run_reads_but_never_rebuilds(self):
        existing = MagicMock()
        existing.get_chunks.return_value = [{"doc_id": "1", "text": "a", "metadata": {}}]
        with patch.object(
            reembed.VectorCreator, "create_vectorstore", return_value=existing
        ) as factory:
            seen, written = reembed.reembed_faiss("s1", batch_size=8, dry_run=True)
        assert (seen, written) == (1, 0)
        assert factory.call_count == 1, "no rebuild store may be constructed"

    def test_empty_index_is_a_no_op(self):
        existing = MagicMock()
        existing.get_chunks.return_value = []
        with patch.object(
            reembed.VectorCreator, "create_vectorstore", return_value=existing
        ):
            assert reembed.reembed_faiss("s1", batch_size=8, dry_run=False) == (0, 0)


class TestPgvectorWithoutTheExtension:
    """Mocked pgvector paths.

    The live tests in ``test_reembed_pgvector_live`` skip wherever the cluster
    has no pgvector build -- which includes CI -- so the SQL shape and the
    batching contract are pinned here too.
    """

    @pytest.fixture
    def store(self):
        """A store whose cursor serves the paginated reads, not a fixed page.

        ``reembed_pgvector`` walks the source by keyset, so a cursor that
        returns the same rows for every ``fetchall`` never terminates. The fake
        answers each of the three queries the function issues from one in-memory
        table, which tests mutate through ``rows``.
        """
        cursor, table = paginating_cursor([(1, "alpha"), (2, "beta"), (3, "gamma")])

        store = MagicMock()
        store._table_name = "documents"
        store._vector_column = "embedding"
        conn = MagicMock()
        conn.cursor.return_value = cursor
        store._get_connection.return_value = conn
        store._embedding.embed_documents.side_effect = lambda texts: [
            [0.5] * 4 for _ in texts
        ]
        with patch.object(
            reembed.VectorCreator, "create_vectorstore", return_value=store
        ), patch.object(reembed.settings, "GRAPHRAG_ENABLED", False):
            yield store, conn, cursor, table

    def test_reads_and_rewrites_every_chunk(self, store):
        _, conn, cursor, _ = store
        seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
        assert (seen, written) == (3, 3)
        cursor.executemany.assert_called_once()
        conn.commit.assert_called()

    def test_pages_are_bounded_by_batch_size(self, store):
        """The read is paginated, so a huge source never lands in memory at once."""
        _, _, cursor, table = store
        table["chunks"] = [(i, f"chunk-{i}") for i in range(1, 11)]
        seen, written = reembed.reembed_pgvector("s1", batch_size=3, dry_run=False)
        assert (seen, written) == (10, 10)
        selects = [
            call for call in cursor.execute.call_args_list
            if "SELECT id, text" in str(call.args[0])
        ]
        # Four pages of at most 3, then one empty page to end the walk.
        assert len(selects) == 5
        assert all(call.args[1][-1] == 3 for call in selects)

    def test_dry_run_counts_without_reading_the_text(self, store):
        fake_store, _, cursor, _ = store
        seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=True)
        assert (seen, written) == (3, 0)
        fake_store._embedding.embed_documents.assert_not_called()
        cursor.executemany.assert_not_called()
        assert any(
            "count(*)" in str(call.args[0]) for call in cursor.execute.call_args_list
        )

    def test_batches_commit_separately(self, store):
        _, conn, cursor, _ = store
        reembed.reembed_pgvector("s1", batch_size=2, dry_run=False)
        # 3 rows at batch 2 is two write transactions.
        assert cursor.executemany.call_count == 2
        assert conn.commit.call_count == 2

    def test_failed_batch_rolls_back_and_raises(self, store):
        _, conn, cursor, _ = store
        cursor.executemany.side_effect = RuntimeError("write failed")
        with pytest.raises(RuntimeError):
            reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
        conn.rollback.assert_called_once()

    def test_connection_is_returned_even_on_failure(self, store):
        fake_store, _, cursor, _ = store
        cursor.executemany.side_effect = RuntimeError("boom")
        with pytest.raises(RuntimeError):
            reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
        fake_store.close.assert_called_once()

    def test_null_text_does_not_crash_the_embed_call(self, store):
        fake_store, _, _, table = store
        table["chunks"] = [(1, None), (2, "beta")]
        seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
        assert (seen, written) == (2, 2)
        assert fake_store._embedding.embed_documents.call_args.args[0] == ["", "beta"]

    def test_empty_source_is_a_no_op(self, store):
        fake_store, _, _, table = store
        table["chunks"] = []
        assert reembed.reembed_pgvector("s1", batch_size=64, dry_run=False) == (0, 0)
        fake_store._embedding.embed_documents.assert_not_called()

    def test_source_discovery_returns_sorted_ids(self):
        store = MagicMock()
        store._table_name = "documents"
        cursor = MagicMock()
        cursor.fetchall.return_value = [("b",), ("a",)]
        conn = MagicMock()
        conn.cursor.return_value = cursor
        store._get_connection.return_value = conn
        with patch.object(
            reembed.VectorCreator, "create_vectorstore", return_value=store
        ):
            assert reembed.list_source_ids("pgvector") == ["b", "a"]


class TestFaissSourceDiscovery:
    def test_source_ids_come_from_index_directories(self):
        storage = MagicMock()
        storage.list_files.return_value = [
            "indexes/src-a/index.faiss",
            "indexes/src-a/index.pkl",
            "indexes/src-b/index.faiss",
        ]
        with patch(
            "docsgpt.storage.storage_creator.StorageCreator.get_storage",
            return_value=storage,
        ):
            assert reembed.list_source_ids("faiss") == ["src-a", "src-b"]

    def test_storage_failure_is_reported_as_a_usable_error(self):
        storage = MagicMock()
        storage.list_files.side_effect = OSError("permission denied")
        with patch(
            "docsgpt.storage.storage_creator.StorageCreator.get_storage",
            return_value=storage,
        ):
            with pytest.raises(reembed.ReembedError, match="permission denied"):
                reembed.list_source_ids("faiss")

    def test_unsupported_store_error_names_the_alternatives(self):
        with patch.object(reembed.settings, "VECTOR_STORE", "milvus", create=True):
            assert reembed.main([]) == 2


class TestGraphNodeReembedding:
    """``graph_nodes.name_embedding`` seeds every graph traversal.

    It is written once at extraction time and never revisited, so rewriting
    only the chunk table leaves the graph seeding from the previous model --
    and since mpnet and granite-311m are both 768-dimensional, the column
    accepts the mismatch and nothing reports it.
    """

    @pytest.fixture
    def graph(self):
        store = MagicMock()
        store._embedding.embed_documents.side_effect = lambda texts: [
            [0.5] * 4 for _ in texts
        ]
        cursor = MagicMock()
        cursor.fetchone.return_value = ("graph_nodes",)
        cursor.fetchall.return_value = [
            ("n1", "Alpha"),
            ("n2", "Beta"),
            ("n3", "Gamma"),
        ]
        conn = MagicMock()
        conn.cursor.return_value = cursor
        return store, conn, cursor

    def test_rewrites_every_node_name(self, graph):
        store, conn, cursor = graph
        written = reembed.reembed_graph_nodes(
            store, conn, "s1", batch_size=64, dry_run=False
        )
        assert written == 3
        store._embedding.embed_documents.assert_called_once_with(
            ["Alpha", "Beta", "Gamma"]
        )
        statement = cursor.executemany.call_args.args[0]
        assert "graph_nodes" in statement and "name_embedding" in statement
        conn.commit.assert_called()

    def test_dry_run_counts_without_embedding(self, graph):
        store, conn, cursor = graph
        assert reembed.reembed_graph_nodes(store, conn, "s1", 64, dry_run=True) == 3
        store._embedding.embed_documents.assert_not_called()
        cursor.executemany.assert_not_called()

    def test_missing_table_is_a_no_op(self, graph):
        store, conn, cursor = graph
        cursor.fetchone.return_value = (None,)
        assert reembed.reembed_graph_nodes(store, conn, "s1", 64, dry_run=False) == 0
        store._embedding.embed_documents.assert_not_called()

    def test_batches_commit_separately(self, graph):
        store, conn, cursor = graph
        reembed.reembed_graph_nodes(store, conn, "s1", batch_size=2, dry_run=False)
        assert cursor.executemany.call_count == 2
        assert conn.commit.call_count == 2

    def test_failed_batch_rolls_back_and_raises(self, graph):
        store, conn, cursor = graph
        cursor.executemany.side_effect = RuntimeError("write failed")
        with pytest.raises(RuntimeError, match="write failed"):
            reembed.reembed_graph_nodes(store, conn, "s1", 64, dry_run=False)
        conn.rollback.assert_called_once()

    def test_pgvector_run_skips_the_graph_when_disabled(self):
        store, conn, cursor = self._pgvector_mocks()
        with patch.object(
            reembed.VectorCreator, "create_vectorstore", return_value=store
        ), patch.object(reembed.settings, "GRAPHRAG_ENABLED", False):
            reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
        assert not self._graph_statements(cursor)

    def test_pgvector_run_reembeds_the_graph_when_enabled(self):
        store, conn, cursor = self._pgvector_mocks()
        with patch.object(
            reembed.VectorCreator, "create_vectorstore", return_value=store
        ), patch.object(reembed.settings, "GRAPHRAG_ENABLED", True):
            reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
        assert self._graph_statements(cursor)

    @staticmethod
    def _pgvector_mocks():
        cursor, _ = paginating_cursor(
            [(1, "alpha"), (2, "beta")],
            graph_rows=[("n1", "Alpha"), ("n2", "Beta")],
        )
        store = MagicMock()
        store._table_name = "documents"
        store._vector_column = "embedding"
        conn = MagicMock()
        conn.cursor.return_value = cursor
        store._get_connection.return_value = conn
        store._embedding.embed_documents.side_effect = lambda texts: [
            [0.5] * 4 for _ in texts
        ]
        return store, conn, cursor

    @staticmethod
    def _graph_statements(cursor):
        return [
            call
            for call in cursor.executemany.call_args_list
            if "graph_nodes" in str(call.args[0])
        ]


class TestEmbedsInProcess:
    """A batch job should not round-trip every chunk through the broker."""

    def test_delegation_is_turned_off_for_the_run(self, monkeypatch):
        from docsgpt.core.settings import settings

        monkeypatch.setattr(settings, "EMBEDDINGS_DELEGATE_TO_WORKER", True, raising=False)
        monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector", raising=False)
        seen = {}
        with patch.object(reembed, "run", side_effect=lambda *a, **k: seen.setdefault(
            "delegating", settings.EMBEDDINGS_DELEGATE_TO_WORKER
        ) or 0):
            assert reembed.main(["--dry-run"]) == 0
        assert seen["delegating"] is False


class TestRecordsTheModel:
    """``sources.model`` is what the boot mismatch check reads."""

    def test_a_re_embedded_source_is_stamped(self, monkeypatch):
        from docsgpt.core.settings import settings

        monkeypatch.setattr(settings, "EMBEDDINGS_NAME", "new/model", raising=False)
        conn = MagicMock()
        session = MagicMock()
        session.__enter__ = MagicMock(return_value=conn)
        session.__exit__ = MagicMock(return_value=False)
        with patch("docsgpt.storage.db.session.db_session", return_value=session):
            reembed.record_source_model("src-1")
        params = conn.execute.call_args.args[1]
        assert params == {"model": "new/model", "id": "src-1"}

    def test_a_dry_run_stamps_nothing(self, monkeypatch):
        with patch.object(reembed, "record_source_model") as record, patch.object(
            reembed, "list_source_ids", return_value=["a"]
        ), patch.object(reembed, "reembed_pgvector", return_value=(3, 0)):
            reembed.run("pgvector", None, 64, True)
        record.assert_not_called()

    def test_a_real_run_stamps_each_source(self, monkeypatch):
        with patch.object(reembed, "record_source_model") as record, patch.object(
            reembed, "list_source_ids", return_value=["a", "b"]
        ), patch.object(reembed, "reembed_pgvector", return_value=(3, 3)):
            reembed.run("pgvector", None, 64, False)
        assert [c.args[0] for c in record.call_args_list] == ["a", "b"]

    def test_a_failed_source_is_not_stamped(self):
        with patch.object(reembed, "record_source_model") as record, patch.object(
            reembed, "list_source_ids", return_value=["a"]
        ), patch.object(reembed, "reembed_pgvector", side_effect=RuntimeError("boom")):
            assert reembed.run("pgvector", None, 64, False) == 1
        record.assert_not_called()


class TestThePinIsResolved:
    """The script must embed with the model the installation is pinned to.

    ``resolve_embeddings_pin`` runs in ``docsgpt.app``, which this script
    never imports. An install pinned in ``app_metadata`` with no
    ``EMBEDDINGS_NAME`` in the environment -- every stock Kubernetes
    deployment, whose manifests carry no embedding config -- would otherwise
    rewrite its whole index with the legacy code default and stamp
    ``sources.model`` to match, creating the cross-model index this script
    exists to repair.
    """

    def test_main_resolves_the_pin_before_reading_the_store(self):
        order = []
        with patch(
            "docsgpt.storage.db.embeddings_pin.resolve_embeddings_pin",
            side_effect=lambda *a, **k: order.append("pin"),
        ), patch.object(reembed.settings, "VECTOR_STORE", "pgvector", create=True), patch.object(
            reembed, "run", side_effect=lambda *a, **k: (order.append("run"), 0)[1]
        ):
            assert reembed.main([]) == 0
        assert order == ["pin", "run"], "the pin must resolve before anything embeds"

    def test_an_unsupported_store_still_resolved_the_pin_first(self):
        with patch(
            "docsgpt.storage.db.embeddings_pin.resolve_embeddings_pin"
        ) as pin, patch.object(reembed.settings, "VECTOR_STORE", "qdrant", create=True):
            assert reembed.main([]) == 2
        pin.assert_called_once()