Files
stack/tests/llm/test_restamp.py
kert 00dd3fb9b0 test(llm): restamp pagination-loop mechanics with a fake engine (F8)
A fake engine/connection records every SQL call and serves fixed
batches, so the loop mechanics are verified directly: `after` advances
to the last id of each batch, the UPDATE runs once per batch with the
planned patches (not once per row), dry_run=True still counts but
issues no UPDATE, and the loop terminates on the first empty batch.
Deleted the tautological test_sql_shapes (asserted substrings of the
constants against themselves).
2026-09-09 21:04:52 -04:00

155 lines
4.4 KiB
Python

"""llm.restamp — metadata-only anchor backfill on already-indexed chunks."""
import json
import os
import pytest
from llm.restamp import plan_patches, restamp
def test_plan_skips_unchanged_and_patches_changed():
rows = [
(
"a",
"Use 99439 in conjunction with 99490; consent",
{
"codes": "99439 99490",
"families": "CCM",
"elements": "activity=consent relation=addon-of",
},
),
(
"b",
"Use 99439 in conjunction with 99490; consent",
{"codes": "99439 99490"},
),
("c", "nothing", {}),
]
patches = plan_patches(rows)
ids = [p[0] for p in patches]
assert ids == ["b", "c"]
assert json.loads(patches[0][1])["families"] == "CCM"
assert json.loads(patches[1][1]) == {"codes": "", "families": "", "elements": ""}
# ── F8: fake-engine pagination-loop mechanics ─────────────────────────
class _FakeResult:
def __init__(self, rows):
self._rows = rows
def fetchall(self):
return self._rows
class _FakeConn:
"""Records every ``execute`` call (SQL text + params) and serves
*batches* — one list of rows per SELECT, in order — then empty lists
forever once exhausted."""
def __init__(self, calls, batches):
self._calls = calls
self._batches = list(batches)
def execute(self, stmt, params=None):
sql = str(stmt)
self._calls.append((sql, params))
if sql.strip().startswith("SELECT"):
rows = self._batches.pop(0) if self._batches else []
return _FakeResult(rows)
return _FakeResult([])
class _FakeBegin:
def __init__(self, conn):
self._conn = conn
def __enter__(self):
return self._conn
def __exit__(self, *exc):
return False
class _FakeEngine:
"""``restamp`` opens a fresh ``with engine.begin() as conn`` per
batch; every call shares one connection so ``calls`` sees the whole
run in order."""
def __init__(self, batches):
self.calls: list[tuple[str, object]] = []
self._conn = _FakeConn(self.calls, batches)
def begin(self):
return _FakeBegin(self._conn)
def _selects(engine):
return [c for c in engine.calls if c[0].strip().startswith("SELECT")]
def _updates(engine):
return [c for c in engine.calls if c[0].strip().startswith("UPDATE")]
def test_loop_advances_after_to_the_last_id_of_each_batch():
batches = [
[("id1", "text about 99490", {})],
[("id2", "text about 99439", {})],
[("id3", "text about G2058", {})],
]
engine = _FakeEngine(batches)
stats = restamp(engine, collection="rules")
# One SELECT per batch plus the final empty one that ends the loop.
afters = [c[1]["after"] for c in _selects(engine)]
assert afters == ["", "id1", "id2", "id3"]
assert stats["scanned"] == 3
def test_update_executed_per_batch_with_the_planned_patches():
batches = [
[("id1", "text about 99490", {})],
[("id2", "text about 99439", {})],
]
engine = _FakeEngine(batches)
stats = restamp(engine, collection="rules")
updates = _updates(engine)
assert len(updates) == 2 # one UPDATE per batch, not per row
assert stats["updated"] == 2
first_patch_ids = [p["id"] for p in updates[0][1]]
assert first_patch_ids == ["id1"]
def test_dry_run_executes_no_update_but_still_counts():
batches = [[("id1", "text about 99490", {})]]
engine = _FakeEngine(batches)
stats = restamp(engine, collection="rules", dry_run=True)
assert stats["scanned"] == 1
assert stats["updated"] == 1
assert _updates(engine) == []
def test_loop_terminates_on_the_first_empty_batch():
engine = _FakeEngine([])
stats = restamp(engine, collection="rules")
assert stats == {
"collection": "rules",
"scanned": 0,
"updated": 0,
"seconds": stats["seconds"],
}
assert len(_selects(engine)) == 1
@pytest.mark.skipif(not os.environ.get("LLM_DB_PASSWORD"), reason="needs pgvector")
def test_restamp_dry_run_scans_live_rules():
from llm import config as llm_config
from llm.index import _engine
cfg = llm_config.load()
engine = _engine(cfg)
stats = restamp(engine, collection="rules", batch=50, dry_run=True)
assert stats["scanned"] > 0