tests/test_docs_reindex_transaction.py

Rama main · réplica de lectura del repositorio original.

"""A failed vector write must not destroy the previously indexed document."""

from datetime import datetime, timezone

import pytest

from dm.docs import index


class _Cursor:
    def __init__(self, connection):
        self.connection = connection

    def execute(self, statement, params=None):
        self.connection.statements.append(statement)
        self.connection.params.append(params)

    def executemany(self, statement, records):
        self.connection.statements.append(statement)
        if self.connection.fail_insert:
            raise RuntimeError("vector insert failed")

    def fetchmany(self, size):
        self.connection.fetch_sizes.append(size)
        return self.connection.rows[:size]

    def close(self):
        self.connection.closed_cursors += 1


class _Connection:
    def __init__(self, rows, *, fail_insert=False):
        self.rows = rows
        self.fail_insert = fail_insert
        self.statements = []
        self.params = []
        self.fetch_sizes = []
        self.closed_cursors = 0
        self.commits = 0
        self.rollbacks = 0
        self.closed = False

    def cursor(self):
        return _Cursor(self)

    def commit(self):
        self.commits += 1

    def rollback(self):
        self.rollbacks += 1

    def close(self):
        self.closed = True


@pytest.mark.parametrize("fail_insert", [False, True])
def test_reindex_replaces_chunks_and_marker_in_one_transaction(monkeypatch, tmp_path, fail_insert):
    source = tmp_path / "document.txt"
    source.write_text("A useful document", encoding="utf-8")
    connection = _Connection(
        [("DOC1", "manual", "Manual", "", str(source), "new-hash", "old-hash")],
        fail_insert=fail_insert,
    )
    monkeypatch.setattr(index, "init_schema", lambda: None)
    monkeypatch.setattr(index, "connect_vec", lambda *, autocommit: connection if autocommit is False else None)
    monkeypatch.setattr(index, "embed", lambda chunks: [[0.1, 0.2] for _ in chunks])
    monkeypatch.setattr(index, "counts", lambda: (1, 1))

    if fail_insert:
        with pytest.raises(RuntimeError, match="vector insert failed"):
            index.reindex(verbose=False)
        assert connection.commits == 0
        assert connection.rollbacks == 1
        assert not any(statement.startswith("UPDATE document") for statement in connection.statements)
    else:
        assert index.reindex(verbose=False) == (1, 1)
        assert connection.commits == 1
        assert connection.rollbacks == 0
        assert any(statement.startswith("UPDATE document") for statement in connection.statements)
    assert connection.closed
    assert connection.fetch_sizes == [index._MAX_DOCUMENT_CATALOG_ROWS + 1]


@pytest.mark.parametrize("embeddings", [[], [[0.1], [0.2]], None])
def test_reindex_rejects_incomplete_embedding_batches_before_deleting_old_chunks(
    monkeypatch, tmp_path, embeddings
):
    source = tmp_path / "document.txt"
    source.write_text("A useful document", encoding="utf-8")
    connection = _Connection(
        [("DOC1", "manual", "Manual", "", str(source), "new-hash", "old-hash")]
    )
    monkeypatch.setattr(index, "init_schema", lambda: None)
    monkeypatch.setattr(index, "connect_vec", lambda *, autocommit: connection)
    monkeypatch.setattr(index, "embed", lambda chunks: embeddings)

    with pytest.raises(RuntimeError, match="unexpected vector count"):
        index.reindex(verbose=False)

    assert connection.commits == 0
    assert connection.rollbacks == 1
    assert not any(statement.startswith("DELETE FROM doc_chunk") for statement in connection.statements)
    assert connection.closed


def test_reindex_allows_nullable_document_title_in_verbose_output(monkeypatch, tmp_path, capsys):
    source = tmp_path / "document.txt"
    source.write_text("A useful document", encoding="utf-8")
    connection = _Connection(
        [("DOC1", "manual", None, "", str(source), "new-hash", "old-hash")]
    )
    monkeypatch.setattr(index, "init_schema", lambda: None)
    monkeypatch.setattr(index, "connect_vec", lambda *, autocommit: connection)
    monkeypatch.setattr(index, "embed", lambda chunks: [[0.1, 0.2] for _ in chunks])
    monkeypatch.setattr(index, "counts", lambda: (1, 1))

    assert index.reindex(verbose=True) == (1, 1)
    assert "DOC1 «»" in capsys.readouterr().out
    assert connection.commits == 1
    assert connection.closed


def test_reindex_persists_a_utc_derived_naive_database_timestamp(monkeypatch, tmp_path):
    class FrozenDateTime(datetime):
        @classmethod
        def now(cls, tz=None):
            assert tz is timezone.utc
            return cls(2026, 10, 13, 9, 30, tzinfo=timezone.utc)

    source = tmp_path / "document.txt"
    source.write_text("A useful document", encoding="utf-8")
    connection = _Connection(
        [("DOC1", "manual", "Manual", "", str(source), "new-hash", "old-hash")]
    )
    monkeypatch.setattr(index, "datetime", FrozenDateTime)
    monkeypatch.setattr(index, "init_schema", lambda: None)
    monkeypatch.setattr(index, "connect_vec", lambda *, autocommit: connection)
    monkeypatch.setattr(index, "embed", lambda chunks: [[0.1, 0.2] for _ in chunks])
    monkeypatch.setattr(index, "counts", lambda: (1, 1))

    assert index.reindex(verbose=False) == (1, 1)

    update_params = next(
        params for params in connection.params
        if isinstance(params, tuple) and len(params) == 4 and params[0] == "new-hash"
    )
    assert update_params[2] == datetime(2026, 10, 13, 9, 30)
    assert update_params[2].tzinfo is None


def test_reindex_rejects_an_oversized_catalog_before_loading_sources(monkeypatch):
    row = ("DOC1", "manual", "Manual", "", "/unused", "new-hash", "old-hash")
    connection = _Connection([row] * (index._MAX_DOCUMENT_CATALOG_ROWS + 1))
    monkeypatch.setattr(index, "init_schema", lambda: None)
    monkeypatch.setattr(index, "connect_vec", lambda *, autocommit: connection)
    monkeypatch.setattr(
        index,
        "read_document_source",
        lambda path: pytest.fail("an oversized catalog must fail before reading document files"),
    )

    with pytest.raises(RuntimeError, match="document catalog exceeds"):
        index.reindex(verbose=False)

    assert connection.commits == 0
    assert connection.closed
    assert connection.closed_cursors == 2


def test_status_bounds_catalog_reads_and_closes_cursor_and_connection(monkeypatch, capsys):
    connection = _Connection(
        [("DOC1", "manual", 1, "", "已索引", "Manual")]
    )
    monkeypatch.setattr(index, "counts", lambda: (1, 1))
    monkeypatch.setattr(index, "connect", lambda: connection)

    index.status()

    assert "DOC1" in capsys.readouterr().out
    assert connection.fetch_sizes == [index._MAX_DOCUMENT_CATALOG_ROWS + 1]
    assert connection.closed_cursors == 1
    assert connection.closed


def test_document_catalog_rejects_malformed_rows():
    connection = _Connection([("DOC1", "manual")])

    with pytest.raises(RuntimeError, match="malformed row"):
        index._read_document_catalog(connection.cursor(), row_width=7)