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).
155 lines
4.4 KiB
Python
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
|