arborist/tests/test_resume.py

"""Resumable ingest: high-water mark in meta lets a stopped ingest rsync forward."""

from __future__ import annotations

from typing import Iterator

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.source import Source
from arborist.store import connect, get_meta, set_meta


class IndexedSource(Source):
    """Test source that exposes an `id` per doc and a start_id filter.

    Imitates the WikipediaSqlDump resume contract.
    """

    source_type = "indexed_test"

    def __init__(self, n_docs: int = 10):
        self.n_docs = n_docs
        self.start_id = 0
        self.last_id = 0

    def iter_documents(self) -> Iterator[Document]:
        for i in range(1, self.n_docs + 1):
            if i <= self.start_id:
                continue
            self.last_id = i
            yield Document(
                uri=f"test://doc/{i}",
                content=f"document number {i} content " * 10,
                source_type=self.source_type,
                title=f"Doc {i}",
                extra={"indexed_test_id": str(i)},
            )


def test_resume_writes_high_water(tmp_path):
    db = tmp_path / "rh.db"
    conn = connect(db)
    try:
        src = IndexedSource(n_docs=10)
        ingest_source(conn, src, batch_size=5, resume=True)
        # high-water now equals the last id seen
        assert get_meta(conn, "source_high_water:indexed_test") == "10"
    finally:
        conn.close()


def test_resume_skips_processed(tmp_path):
    db = tmp_path / "skip.db"
    conn = connect(db)
    try:
        # Pre-seed a high-water as if a prior run got through doc 4.
        with __import__("arborist").store.transaction(conn):
            set_meta(conn, "source_high_water:indexed_test", "4")
        src = IndexedSource(n_docs=10)
        result = ingest_source(conn, src, batch_size=5, resume=True)
        # Source should have skipped docs 1..4 entirely (start_id=4, then >4 emitted).
        # Note: `seen` counts what the source yielded post-skip.
        assert result.seen == 6
        assert result.inserted == 6
        # high-water now updated to 10
        assert get_meta(conn, "source_high_water:indexed_test") == "10"
    finally:
        conn.close()


def test_resume_idempotent(tmp_path):
    """Running ingest twice with resume yields no new docs the second time."""
    db = tmp_path / "idem.db"
    conn = connect(db)
    try:
        ingest_source(conn, IndexedSource(n_docs=10), batch_size=3, resume=True)
        n1 = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
        ingest_source(conn, IndexedSource(n_docs=10), batch_size=3, resume=True)
        n2 = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
        assert n1 == n2 == 10
    finally:
        conn.close()


def test_resume_continues_after_kill(tmp_path):
    """Simulate a kill mid-batch: high-water is at the last successfully
    flushed batch, restart picks up from there without dups or loss."""
    db = tmp_path / "kill.db"
    conn = connect(db)
    try:
        # First run hits 7 of 10, then "killed" — but with batch_size=5 the
        # high-water should be 5 (one full batch flushed).
        ingest_source(conn, IndexedSource(n_docs=5), batch_size=5, resume=True)
        assert get_meta(conn, "source_high_water:indexed_test") == "5"

        # Resume with the full source — should pick up at 6..10.
        result = ingest_source(conn, IndexedSource(n_docs=10), batch_size=5, resume=True)
        assert result.seen == 5
        assert result.inserted == 5
        n_total = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
        assert n_total == 10
    finally:
        conn.close()