997 lines
38 KiB
Python
997 lines
38 KiB
Python
"""tools/repair_foreign_links.py: the reviewed repair of the #247 dataset links.
|
|
|
|
The fixture is a library whose foreign (no PromptManager node) dataset run
|
|
wrote shared-prefix files that the old scan linked to the wrong branch text.
|
|
Recorded history names each file's saver, so the ladder (``recorded-run``)
|
|
names the right branch.
|
|
"""
|
|
|
|
import contextlib
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import os
|
|
import shutil
|
|
import sqlite3
|
|
import tempfile
|
|
import unittest
|
|
import unittest.mock
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
try:
|
|
from tests.open_handles import assert_closed
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from open_handles import assert_closed
|
|
|
|
from database.models import PromptModel
|
|
from database.operations import PromptDatabase
|
|
from tests import saver_anchor_graphs as shapes
|
|
from tests.test_run_attribution import RecordedRuns
|
|
from tools import repair_foreign_links as tool
|
|
from utils.file_prompt import Decision
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
TEXTS = shapes.BRANCH_TEXTS
|
|
NAME = "IMG_GEN_{:05}_.png"
|
|
|
|
|
|
def png(path, graph):
|
|
info = PngInfo()
|
|
info.add_text("prompt", json.dumps(graph))
|
|
Image.new("RGB", (4, 4)).save(path, pnginfo=info)
|
|
|
|
|
|
def sha(text):
|
|
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def file_sha(path):
|
|
return hashlib.sha256(Path(path).read_bytes()).hexdigest()
|
|
|
|
|
|
def dump(path):
|
|
"""The database's logical content (independent of WAL checkpoints)."""
|
|
conn = sqlite3.connect(path)
|
|
try:
|
|
return list(conn.iterdump())
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def not_running(_comfyui_dir, _port):
|
|
return []
|
|
|
|
|
|
def no_handles(_database):
|
|
return []
|
|
|
|
|
|
class RepairFixture(RecordedRuns, unittest.TestCase):
|
|
"""Rows: A relink to a new prompt, B relink to a stored prompt (with a
|
|
pair link), C keep, D file gone, E two savers (skip), F the UNIQUE case
|
|
(a wrong row beside a correct row), H relink from a rated prompt, G a
|
|
PromptManager run (not a candidate)."""
|
|
|
|
def setUp(self):
|
|
self.tmp = tempfile.mkdtemp()
|
|
self.addCleanup(shutil.rmtree, self.tmp, True)
|
|
self.addCleanup(assert_closed, self, self.tmp)
|
|
self.addCleanup(PromptModel.reset_schema_cache)
|
|
self.output = os.path.join(self.tmp, "output")
|
|
os.makedirs(os.path.join(self.output, "dataset"))
|
|
self.db_path = os.path.join(self.tmp, "prompts.db")
|
|
self.plan_path = os.path.join(self.tmp, "plan.json")
|
|
self.backups = os.path.join(self.tmp, "backups")
|
|
os.makedirs(self.backups)
|
|
self.comfy = os.path.join(self.tmp, "comfy")
|
|
os.makedirs(self.comfy)
|
|
Path(self.comfy, "main.py").write_text("# ComfyUI\n", encoding="utf-8")
|
|
self.files = {}
|
|
for key, number in zip("ABCDEFHG", range(1, 9)):
|
|
path = os.path.join(self.output, "dataset", NAME.format(number))
|
|
png(path, shapes.shared_prefix())
|
|
self.files[key] = path
|
|
db = PromptDatabase(self.db_path)
|
|
try:
|
|
self.build(db)
|
|
finally:
|
|
db.close_all()
|
|
os.remove(self.files["D"])
|
|
|
|
def build(self, db):
|
|
save = lambda text, **kw: db.save_prompt( # noqa: E731
|
|
text=text, prompt_hash=generate_prompt_hash(text), **kw
|
|
)
|
|
self.p0 = save(TEXTS[0], category="scanned", tags=["auto-scanned"])
|
|
self.p2 = save(TEXTS[2], category="scanned", tags=["auto-scanned"])
|
|
self.orphan = save(
|
|
"an old guess",
|
|
category="scanned",
|
|
tags=["auto-scanned"],
|
|
notes="Auto-scanned from IMG_GEN_00001_.png",
|
|
)
|
|
self.rated = save("a rated guess", rating=5)
|
|
link = lambda prompt, key: db.link_image_to_prompt( # noqa: E731
|
|
prompt, self.files[key], metadata={}
|
|
)
|
|
self.rows = {
|
|
"A": link(self.orphan, "A"),
|
|
"B": link(self.p0, "B"),
|
|
"C": link(self.p0, "C"),
|
|
"D": link(self.p0, "D"),
|
|
"E": link(self.p0, "E"),
|
|
"F": link(self.p0, "F"),
|
|
"F_ok": link(self.p2, "F"),
|
|
"H": link(self.rated, "H"),
|
|
"G": link(self.p0, "G"),
|
|
}
|
|
conn = db.model.get_connection()
|
|
with conn:
|
|
pair = conn.execute(
|
|
"INSERT INTO prompt_pairs(positive_prompt_id) VALUES(?)", (self.p0,)
|
|
).lastrowid
|
|
conn.execute(
|
|
"INSERT INTO image_prompt_pairs(image_id, pair_id, provenance) "
|
|
"VALUES(?,?,'{}')",
|
|
(self.rows["B"], pair),
|
|
)
|
|
name = lambda key: ("dataset", os.path.basename(self.files[key])) # noqa
|
|
# Savers of the shared-prefix graph: "13", "23", "33", "43" (branches 1-4)
|
|
self.record_run(
|
|
db,
|
|
self.output,
|
|
[name(k) for k in "CDEAEBFH"],
|
|
[],
|
|
nodes=["13", "13", "13", "23", "23", "33", "33", "43"],
|
|
)
|
|
self.record_run(db, self.output, [name("G")], [TEXTS[0]], item=1, nodes=["23"])
|
|
|
|
def dry_run(self, *args):
|
|
out = io.StringIO()
|
|
code = tool.main(
|
|
[
|
|
self.db_path,
|
|
"--output-dir",
|
|
self.output,
|
|
"--plan",
|
|
self.plan_path,
|
|
*args,
|
|
],
|
|
out=out,
|
|
)
|
|
return code, out.getvalue()
|
|
|
|
def apply(self, *args, probe=not_running, handles=no_handles):
|
|
out = io.StringIO()
|
|
code = tool.main(
|
|
[
|
|
self.db_path,
|
|
"--apply",
|
|
self.plan_path,
|
|
"--comfyui-dir",
|
|
os.path.join(self.tmp, "comfy"),
|
|
"--backup-dir",
|
|
self.backups,
|
|
*args,
|
|
],
|
|
out=out,
|
|
probe=probe,
|
|
handles=handles,
|
|
)
|
|
return code, out.getvalue()
|
|
|
|
def plan(self):
|
|
with open(self.plan_path, encoding="utf-8") as stream:
|
|
return json.load(stream)
|
|
|
|
def write_plan(self, plan):
|
|
with open(self.plan_path, "w", encoding="utf-8") as stream:
|
|
json.dump(plan, stream)
|
|
|
|
def query(self, sql, args=()):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
return conn.execute(sql, args).fetchall()
|
|
finally:
|
|
conn.close()
|
|
|
|
def query_write(self, sql, args):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
with conn:
|
|
conn.execute(sql, args)
|
|
finally:
|
|
conn.close()
|
|
|
|
def links(self):
|
|
return dict(self.query("SELECT id, prompt_id FROM generated_images"))
|
|
|
|
|
|
class DryRunTests(RepairFixture):
|
|
def test_plan_names_relink_keep_and_skip_rows(self):
|
|
code, out = self.dry_run()
|
|
self.assertEqual(code, 0, out)
|
|
items = {item["row_id"]: item for item in self.plan()["rows"]}
|
|
rows = self.rows
|
|
self.assertNotIn(rows["G"], items, "a PromptManager run is not foreign")
|
|
actions = {key: items[rows[key]]["action"] for key in rows if key != "G"}
|
|
self.assertEqual(
|
|
actions,
|
|
{
|
|
"A": "relink",
|
|
"B": "relink",
|
|
"C": "keep",
|
|
"D": "skip",
|
|
"E": "skip",
|
|
"F": "relink",
|
|
"F_ok": "keep",
|
|
"H": "relink",
|
|
},
|
|
)
|
|
a = items[rows["A"]]
|
|
self.assertEqual(a["current_prompt_id"], self.orphan)
|
|
self.assertEqual(a["current_text_sha256"], sha("an old guess"))
|
|
self.assertEqual(a["target_text"], TEXTS[1])
|
|
self.assertEqual(a["target_hash"], generate_prompt_hash(TEXTS[1]))
|
|
self.assertEqual(a["target_prompt_id"], "new")
|
|
self.assertEqual(a["file"], self.files["A"])
|
|
self.assertEqual(items[rows["B"]]["target_prompt_id"], self.p2)
|
|
self.assertEqual(items[rows["H"]]["target_prompt_id"], "new")
|
|
self.assertEqual(items[rows["D"]]["reason"], "file gone")
|
|
self.assertEqual(
|
|
items[rows["E"]]["reason"], "not decided by the recorded saver"
|
|
)
|
|
self.assertIn("relink: 4", out)
|
|
self.assertIn("keep: 2", out)
|
|
self.assertIn("skip: 2", out)
|
|
|
|
def test_orphan_audit_lists_the_prompts_losing_every_link(self):
|
|
self.dry_run()
|
|
plan = self.plan()
|
|
orphans = {o["id"]: o for o in plan["orphans"]}
|
|
self.assertEqual(set(orphans), {self.orphan, self.rated})
|
|
self.assertEqual(orphans[self.orphan]["blockers"], [])
|
|
self.assertTrue(orphans[self.orphan]["safe"])
|
|
self.assertEqual(orphans[self.rated]["blockers"], ["usage or rating"])
|
|
self.assertFalse(orphans[self.rated]["safe"])
|
|
self.assertFalse(plan["delete_orphans"], "off by default")
|
|
self.dry_run("--delete-orphans")
|
|
self.assertTrue(self.plan()["delete_orphans"])
|
|
|
|
def test_dry_run_writes_nothing_and_writes_the_markdown_list(self):
|
|
before = file_sha(self.db_path)
|
|
listing = sorted(os.listdir(self.tmp))
|
|
code, out = self.dry_run()
|
|
self.assertEqual(code, 0)
|
|
self.assertEqual(file_sha(self.db_path), before)
|
|
self.assertIn("Dry run: nothing was changed.", out)
|
|
# A read-only reader of a WAL library may leave its empty side files
|
|
sidecars = {"prompts.db-shm", "prompts.db-wal"}
|
|
self.assertEqual(
|
|
sorted(set(os.listdir(self.tmp)) - sidecars),
|
|
sorted(set(listing + ["plan.json", "plan.list.md"]) - sidecars),
|
|
)
|
|
wal = Path(self.tmp, "prompts.db-wal")
|
|
self.assertEqual(wal.stat().st_size if wal.exists() else 0, 0)
|
|
markdown = Path(self.tmp, "plan.list.md").read_text(encoding="utf-8")
|
|
self.assertIn(f"| {self.rows['A']} |", markdown)
|
|
long_path = os.path.join("dataset", NAME.format(1))
|
|
self.assertIn(f"| {long_path} |", markdown, "file paths are never cut")
|
|
self.assertIn(generate_prompt_hash(TEXTS[1])[:12], markdown)
|
|
self.assertIn("new", markdown)
|
|
self.assertIn("would lose every link", markdown)
|
|
|
|
|
|
class ManualLinkTests(RepairFixture):
|
|
"""A prompt the reviewer named in Needs review (#536) is not ours to undo."""
|
|
|
|
def mark_manual(self, key):
|
|
row = self.rows[key]
|
|
((prompt,),) = self.query(
|
|
"SELECT prompt_id FROM generated_images WHERE id=?", (row,)
|
|
)
|
|
self.query_write(
|
|
"INSERT OR IGNORE INTO prompt_pairs(positive_prompt_id,created_at,"
|
|
"updated_at) VALUES(?,'2026-10-30T00:00:00Z','2026-10-30T00:00:00Z')",
|
|
(prompt,),
|
|
)
|
|
((pair,),) = self.query(
|
|
"SELECT id FROM prompt_pairs WHERE positive_prompt_id=? "
|
|
"AND negative_prompt_id IS NULL",
|
|
(prompt,),
|
|
)
|
|
self.query_write(
|
|
"INSERT INTO image_prompt_pairs(pair_id,image_id,provenance,created_at) "
|
|
"VALUES(?,?,?,'2026-10-30T00:00:00Z')",
|
|
(pair, row, '{"kind":"manual-link"}'),
|
|
)
|
|
|
|
def test_a_manual_link_is_skipped_in_the_plan(self):
|
|
self.mark_manual("A")
|
|
code, out = self.dry_run()
|
|
self.assertEqual(code, 0, out)
|
|
items = {item["row_id"]: item for item in self.plan()["rows"]}
|
|
item = items[self.rows["A"]]
|
|
self.assertEqual(item["action"], "skip")
|
|
self.assertEqual(item["reason"], "linked by hand in Needs review")
|
|
self.assertEqual(items[self.rows["B"]]["action"], "relink")
|
|
|
|
def test_a_manual_link_made_after_the_plan_refuses_the_apply(self):
|
|
code, out = self.dry_run()
|
|
self.assertEqual(code, 0, out)
|
|
self.mark_manual("A")
|
|
before = self.links()
|
|
code, out = self.apply()
|
|
self.assertNotEqual(code, 0, out)
|
|
self.assertIn("action changed", out)
|
|
self.assertEqual(self.links(), before)
|
|
|
|
|
|
class DryRunReviewTests(RepairFixture):
|
|
def test_plan_records_its_library_and_fingerprint(self):
|
|
self.dry_run()
|
|
plan = self.plan()
|
|
self.assertEqual(
|
|
plan["source_db"], os.path.normcase(os.path.realpath(self.db_path))
|
|
)
|
|
self.assertEqual(plan["max_row_id"], max(self.links()))
|
|
|
|
def test_a_markdown_plan_path_is_refused(self):
|
|
out = io.StringIO()
|
|
with (
|
|
self.assertRaises(SystemExit),
|
|
contextlib.redirect_stderr(io.StringIO()),
|
|
):
|
|
tool.main(
|
|
[
|
|
self.db_path,
|
|
"--output-dir",
|
|
self.output,
|
|
"--plan",
|
|
os.path.join(self.tmp, "plan.md"),
|
|
],
|
|
out=out,
|
|
)
|
|
self.assertFalse(os.path.exists(os.path.join(self.tmp, "plan.md")))
|
|
|
|
def test_keep_compares_the_linked_text_not_the_stored_hash(self):
|
|
self.query_write("UPDATE prompts SET hash='stale' WHERE id=?", (self.p0,))
|
|
self.dry_run()
|
|
items = {item["row_id"]: item for item in self.plan()["rows"]}
|
|
self.assertEqual(items[self.rows["C"]]["action"], "keep")
|
|
|
|
def test_delete_orphans_is_rehearsed_with_the_deletes(self):
|
|
code, out = self.dry_run("--delete-orphans")
|
|
self.assertEqual(code, 0, out)
|
|
simulated = self.plan()["simulated"]
|
|
self.assertEqual(simulated["result"]["orphans_deleted"], 1)
|
|
self.assertEqual(
|
|
simulated["after"]["prompts"], simulated["before"]["prompts"] + 2 - 1
|
|
)
|
|
|
|
def test_a_file_on_another_drive_is_listed_by_its_full_path(self):
|
|
self.dry_run()
|
|
plan = self.plan()
|
|
with unittest.mock.patch.object(
|
|
tool.os.path, "relpath", side_effect=ValueError("other drive")
|
|
):
|
|
listing = tool.markdown(plan, {})
|
|
self.assertIn(self.files["A"], listing)
|
|
|
|
|
|
class ApplyTests(RepairFixture):
|
|
def setUp(self):
|
|
super().setUp()
|
|
code, out = self.dry_run()
|
|
self.assertEqual(code, 0, out)
|
|
|
|
def test_apply_relinks_creates_targets_and_drops_pair_links(self):
|
|
code, out = self.apply()
|
|
self.assertEqual(code, 0, out)
|
|
links = self.links()
|
|
new = dict(self.query("SELECT hash, id FROM prompts"))
|
|
target_a = new[generate_prompt_hash(TEXTS[1])]
|
|
target_h = new[generate_prompt_hash(TEXTS[3])]
|
|
self.assertEqual(links[self.rows["A"]], target_a)
|
|
self.assertEqual(links[self.rows["B"]], self.p2)
|
|
self.assertEqual(links[self.rows["H"]], target_h)
|
|
self.assertEqual(links[self.rows["C"]], self.p0, "keep")
|
|
self.assertEqual(links[self.rows["E"]], self.p0, "skip")
|
|
self.assertEqual(links[self.rows["G"]], self.p0, "not foreign")
|
|
self.assertEqual(
|
|
self.query(
|
|
"SELECT COUNT(*) FROM image_prompt_pairs WHERE image_id=?",
|
|
(self.rows["B"],),
|
|
),
|
|
[(0,)],
|
|
)
|
|
category, notes = self.query(
|
|
"SELECT category, notes FROM prompts WHERE id=?", (target_a,)
|
|
)[0]
|
|
self.assertEqual(category, "scanned")
|
|
self.assertEqual(
|
|
notes, f"Auto-scanned from {os.path.basename(self.files['A'])}"
|
|
)
|
|
tags = self.query(
|
|
"SELECT t.name FROM prompt_tags pt JOIN tags t ON t.id=pt.tag_id "
|
|
"WHERE pt.prompt_id=?",
|
|
(target_a,),
|
|
)
|
|
self.assertEqual(tags, [("auto-scanned",)])
|
|
self.assertIn("integrity_check", out)
|
|
self.assertIn("foreign_key_check", out)
|
|
self.assertRegex(out, r"integrity_check\s+ok\s+ok")
|
|
# Orphans are kept unless the plan chose to delete them
|
|
self.assertTrue(self.query("SELECT 1 FROM prompts WHERE id=?", (self.orphan,)))
|
|
|
|
def test_unique_case_deletes_the_wrong_row(self):
|
|
code, out = self.apply()
|
|
self.assertEqual(code, 0, out)
|
|
links = self.links()
|
|
self.assertNotIn(self.rows["F"], links)
|
|
self.assertEqual(links[self.rows["F_ok"]], self.p2)
|
|
rows = self.query(
|
|
"SELECT prompt_id FROM generated_images WHERE image_path=?",
|
|
(self.files["F"],),
|
|
)
|
|
self.assertEqual(rows, [(self.p2,)])
|
|
|
|
def test_backup_is_taken_before_writing(self):
|
|
before = self.links()
|
|
code, out = self.apply()
|
|
self.assertEqual(code, 0, out)
|
|
names = os.listdir(self.backups)
|
|
self.assertEqual(len(names), 1)
|
|
self.assertRegex(names[0], r"^prompts\.db\.pre-247-repair_\d{8}_\d{6}$")
|
|
backup = os.path.join(self.backups, names[0])
|
|
conn = sqlite3.connect(backup)
|
|
try:
|
|
saved = dict(conn.execute("SELECT id, prompt_id FROM generated_images"))
|
|
finally:
|
|
conn.close()
|
|
self.assertEqual(saved, before)
|
|
self.assertIn(f"Backup: {backup}", out)
|
|
|
|
def test_backup_beside_the_database_by_default(self):
|
|
out = io.StringIO()
|
|
code = tool.main(
|
|
[
|
|
self.db_path,
|
|
"--apply",
|
|
self.plan_path,
|
|
"--comfyui-dir",
|
|
os.path.join(self.tmp, "comfy"),
|
|
],
|
|
out=out,
|
|
probe=not_running,
|
|
handles=no_handles,
|
|
)
|
|
self.assertEqual(code, 0, out.getvalue())
|
|
self.assertTrue(
|
|
any(
|
|
n.startswith("prompts.db.pre-247-repair_") for n in os.listdir(self.tmp)
|
|
)
|
|
)
|
|
|
|
def test_refuses_while_comfyui_runs(self):
|
|
before = file_sha(self.db_path)
|
|
code, out = self.apply(probe=lambda _d, _p: ["process 42 (main.py)"])
|
|
self.assertEqual(code, 2)
|
|
self.assertIn("REFUSED", out)
|
|
self.assertIn("process 42", out)
|
|
self.assertEqual(os.listdir(self.backups), [], "no backup either")
|
|
self.assertEqual(file_sha(self.db_path), before)
|
|
|
|
def test_rehearsal_skips_the_running_check_on_a_copy_only(self):
|
|
code, out = self.apply("--rehearsal", probe=lambda _d, _p: ["port 8188"])
|
|
self.assertEqual(code, 0, out)
|
|
self.assertIn("rehearsal", out)
|
|
# The live library of that ComfyUI folder: --rehearsal is refused
|
|
live_dir = os.path.join(self.tmp, "comfy", "user", "__promptmanager")
|
|
os.makedirs(live_dir)
|
|
live = os.path.join(live_dir, "prompts.db")
|
|
shutil.copy(self.db_path, live)
|
|
out = io.StringIO()
|
|
code = tool.main(
|
|
[
|
|
live,
|
|
"--apply",
|
|
self.plan_path,
|
|
"--comfyui-dir",
|
|
os.path.join(self.tmp, "comfy"),
|
|
"--rehearsal",
|
|
],
|
|
out=out,
|
|
probe=lambda _d, _p: [],
|
|
handles=no_handles,
|
|
)
|
|
self.assertEqual(code, 2)
|
|
self.assertIn("--rehearsal", out.getvalue())
|
|
|
|
def test_drift_refuses_and_changes_nothing(self):
|
|
cases = {
|
|
"row moved": lambda: self.query_write(
|
|
"UPDATE generated_images SET prompt_id=? WHERE id=?",
|
|
(self.p2, self.rows["H"]),
|
|
),
|
|
"text changed": lambda: self.query_write(
|
|
"UPDATE prompts SET text='an older guess' WHERE id=?", (self.orphan,)
|
|
),
|
|
"file changed": lambda: png(self.files["B"], shapes.branch_bench()),
|
|
"target stored since": lambda: self.query_write(
|
|
"INSERT INTO prompts(text, hash) VALUES(?,?)",
|
|
(TEXTS[1], generate_prompt_hash(TEXTS[1])),
|
|
),
|
|
}
|
|
pristine = os.path.join(self.tmp, "pristine.db")
|
|
shutil.copy(self.db_path, pristine)
|
|
image = Path(self.files["B"]).read_bytes()
|
|
for name, drift in cases.items():
|
|
with self.subTest(name):
|
|
shutil.copy(pristine, self.db_path)
|
|
Path(self.files["B"]).write_bytes(image)
|
|
shutil.rmtree(self.backups)
|
|
os.makedirs(self.backups)
|
|
drift()
|
|
before = dump(self.db_path)
|
|
code, out = self.apply()
|
|
self.assertEqual(code, 2, out)
|
|
self.assertIn("REFUSED (rolled back, nothing changed)", out)
|
|
self.assertEqual(dump(self.db_path), before)
|
|
|
|
def test_late_failure_rolls_back_every_write(self):
|
|
plan = self.plan()
|
|
plan["delete_orphans"] = True
|
|
self.write_plan(plan)
|
|
# The safe orphan gained a user tag after the dry run
|
|
self.query_write("INSERT INTO tags(name) VALUES('mine')", ())
|
|
self.query_write(
|
|
"INSERT INTO prompt_tags(prompt_id, tag_id) "
|
|
"SELECT ?, id FROM tags WHERE name='mine'",
|
|
(self.orphan,),
|
|
)
|
|
before = self.links()
|
|
prompts = self.query("SELECT COUNT(*) FROM prompts")
|
|
code, out = self.apply()
|
|
self.assertEqual(code, 2, out)
|
|
self.assertIn(f"prompt {self.orphan}", out)
|
|
self.assertIn("user tags", out)
|
|
self.assertEqual(self.links(), before, "the relinks were rolled back")
|
|
self.assertEqual(self.query("SELECT COUNT(*) FROM prompts"), prompts)
|
|
|
|
def test_orphans_are_deleted_only_when_chosen_and_unreferenced(self):
|
|
plan = self.plan()
|
|
plan["delete_orphans"] = True
|
|
self.write_plan(plan)
|
|
code, out = self.apply()
|
|
self.assertEqual(code, 0, out)
|
|
self.assertFalse(self.query("SELECT 1 FROM prompts WHERE id=?", (self.orphan,)))
|
|
self.assertTrue(
|
|
self.query("SELECT 1 FROM prompts WHERE id=?", (self.rated,)),
|
|
"a prompt with a rating is never deleted",
|
|
)
|
|
self.assertIn("orphans_deleted", out)
|
|
|
|
def test_an_applied_plan_refuses_a_second_time(self):
|
|
self.assertEqual(self.apply()[0], 0)
|
|
before = dump(self.db_path)
|
|
code, out = self.apply()
|
|
self.assertEqual(code, 2)
|
|
self.assertIn("REFUSED", out)
|
|
self.assertEqual(dump(self.db_path), before)
|
|
|
|
|
|
class ApplySafetyTests(RepairFixture):
|
|
def setUp(self):
|
|
super().setUp()
|
|
code, out = self.dry_run()
|
|
self.assertEqual(code, 0, out)
|
|
|
|
def assert_unchanged(self, before, code, out):
|
|
self.assertEqual(code, 2, out)
|
|
self.assertEqual(dump(self.db_path), before)
|
|
|
|
def test_refuses_when_another_process_has_the_library_open(self):
|
|
before = dump(self.db_path)
|
|
held = lambda _db: ["process 7 has prompts.db-wal open"] # noqa: E731
|
|
for args in ((), ("--rehearsal",)):
|
|
with self.subTest(args):
|
|
code, out = self.apply(*args, handles=held)
|
|
self.assert_unchanged(before, code, out)
|
|
self.assertIn("process 7", out)
|
|
self.assertEqual(os.listdir(self.backups), [])
|
|
|
|
def test_the_running_check_runs_again_just_before_writing(self):
|
|
calls = []
|
|
|
|
def late(_d, _p):
|
|
calls.append(1)
|
|
return [] if len(calls) == 1 else ["something listens on port 8188"]
|
|
|
|
before = dump(self.db_path)
|
|
code, out = self.apply(probe=late)
|
|
self.assert_unchanged(before, code, out)
|
|
self.assertEqual(len(calls), 2)
|
|
self.assertIn("rolled back", out)
|
|
self.assertIn("Backup:", out)
|
|
|
|
def test_without_proc_only_an_explicit_stop_lets_it_through(self):
|
|
blind = lambda _db: [tool.NO_PROC] # noqa: E731
|
|
before = dump(self.db_path)
|
|
code, out = self.apply(handles=blind)
|
|
self.assert_unchanged(before, code, out)
|
|
self.assertIn("--i-stopped-comfyui", out)
|
|
code, out = self.apply("--i-stopped-comfyui", handles=blind)
|
|
self.assertEqual(code, 0, out)
|
|
|
|
def test_a_folder_without_main_py_is_refused(self):
|
|
os.remove(os.path.join(self.comfy, "main.py"))
|
|
before = dump(self.db_path)
|
|
code, out = self.apply()
|
|
self.assert_unchanged(before, code, out)
|
|
self.assertIn("main.py", out)
|
|
|
|
def test_a_plan_for_another_library_is_refused(self):
|
|
other = os.path.join(self.tmp, "other.db")
|
|
shutil.copy(self.db_path, other)
|
|
out = io.StringIO()
|
|
code = tool.main(
|
|
[other, "--apply", self.plan_path, "--comfyui-dir", self.comfy],
|
|
out=out,
|
|
probe=not_running,
|
|
handles=no_handles,
|
|
)
|
|
self.assertEqual(code, 2)
|
|
self.assertIn("made for", out.getvalue())
|
|
|
|
def test_opening_the_library_must_not_change_its_schema(self):
|
|
self.query_write("DROP INDEX idx_prompt_pairs_positive_only", ())
|
|
PromptModel.reset_schema_cache()
|
|
code, out = self.apply()
|
|
self.assertEqual(code, 2, out)
|
|
self.assertIn("schema", out)
|
|
self.assertIn("live code version", out)
|
|
|
|
def test_an_existing_backup_file_is_refused(self):
|
|
stamp = "20261028_000000"
|
|
Path(self.backups, f"prompts.db.pre-247-repair_{stamp}").write_bytes(b"x")
|
|
before = dump(self.db_path)
|
|
with unittest.mock.patch.object(tool, "_stamp", return_value=stamp):
|
|
code, out = self.apply()
|
|
self.assert_unchanged(before, code, out)
|
|
self.assertIn("already exists", out)
|
|
|
|
def test_a_backup_that_does_not_match_is_refused(self):
|
|
real = tool.backup_database
|
|
|
|
def short(source, target):
|
|
real(source, target)
|
|
conn = sqlite3.connect(target)
|
|
try:
|
|
with conn:
|
|
conn.execute("DELETE FROM image_prompt_pairs")
|
|
finally:
|
|
conn.close()
|
|
|
|
before = dump(self.db_path)
|
|
with unittest.mock.patch.object(tool, "backup_database", short):
|
|
code, out = self.apply()
|
|
self.assert_unchanged(before, code, out)
|
|
self.assertIn("backup", out)
|
|
|
|
def test_a_changed_target_text_is_refused(self):
|
|
plan = self.plan()
|
|
for item in plan["rows"]:
|
|
if item["row_id"] == self.rows["A"]:
|
|
item["target_text"] = "something else"
|
|
self.write_plan(plan)
|
|
before = dump(self.db_path)
|
|
code, out = self.apply()
|
|
self.assert_unchanged(before, code, out)
|
|
self.assertIn("target_text", out)
|
|
|
|
def test_an_orphan_gone_or_changed_at_apply_is_refused(self):
|
|
plan = self.plan()
|
|
plan["delete_orphans"] = True
|
|
plan["orphans"][0]["text_sha256"] = "0" * 64
|
|
self.write_plan(plan)
|
|
before = dump(self.db_path)
|
|
code, out = self.apply()
|
|
self.assert_unchanged(before, code, out)
|
|
self.assertIn("gone or text changed", out)
|
|
|
|
def test_any_error_rolls_back_and_names_the_backup(self):
|
|
before = dump(self.db_path)
|
|
with unittest.mock.patch.object(
|
|
tool, "_relink", side_effect=RuntimeError("disk on fire")
|
|
):
|
|
code, out = self.apply()
|
|
self.assert_unchanged(before, code, out)
|
|
self.assertIn("rolled back", out)
|
|
self.assertIn("disk on fire", out)
|
|
self.assertIn(self.backups, out)
|
|
|
|
def test_a_failed_integrity_check_after_writing_exits_1(self):
|
|
real = tool.counts
|
|
seen = []
|
|
|
|
def counts(path, readonly=True):
|
|
found = real(path, readonly)
|
|
if path == self.db_path:
|
|
seen.append(path)
|
|
if len(seen) > 1:
|
|
return {**found, "integrity_check": "corrupt"}
|
|
return found
|
|
|
|
with unittest.mock.patch.object(tool, "counts", counts):
|
|
code, out = self.apply()
|
|
self.assertEqual(code, 1, out)
|
|
self.assertIn("restore", out)
|
|
|
|
|
|
class BlockerTests(RepairFixture):
|
|
"""Each reference that keeps a prompt from being deleted."""
|
|
|
|
def test_every_blocker(self):
|
|
pid = self.orphan
|
|
cases = {
|
|
"linked images": None,
|
|
"saved pairs": "INSERT INTO prompt_pairs(positive_prompt_id) VALUES(:id)",
|
|
"recorded runs": "UPDATE execution_run_contexts "
|
|
"SET positive_prompt_id=:id",
|
|
"library state": "INSERT INTO prompt_library_state(prompt_id) "
|
|
"VALUES(:id)",
|
|
"user tags": "INSERT INTO prompt_tags(prompt_id, tag_id) "
|
|
"SELECT :id, id FROM tags WHERE name='mine'",
|
|
"usage or rating": "UPDATE prompts SET run_count=3 WHERE id=:id",
|
|
"category": "UPDATE prompts SET category='portraits' WHERE id=:id",
|
|
"notes": "UPDATE prompts SET notes='mine' WHERE id=:id",
|
|
}
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute("INSERT INTO tags(name) VALUES('mine')")
|
|
conn.commit()
|
|
for name, sql in cases.items():
|
|
with self.subTest(name):
|
|
conn.execute("SAVEPOINT blocker")
|
|
if sql is None:
|
|
unused = {pid}
|
|
else:
|
|
conn.execute(
|
|
"DELETE FROM generated_images WHERE prompt_id=?", (pid,)
|
|
)
|
|
conn.execute(sql, {"id": pid})
|
|
unused = {pid}
|
|
found = tool.blockers(conn, pid, unused)
|
|
self.assertTrue(any(item.startswith(name) for item in found), found)
|
|
conn.execute("ROLLBACK TO blocker")
|
|
conn.execute("RELEASE blocker")
|
|
conn.execute("DELETE FROM generated_images WHERE prompt_id=?", (pid,))
|
|
self.assertEqual(tool.blockers(conn, pid, {pid}), [])
|
|
self.assertEqual(
|
|
tool.blockers(conn, pid, set()), ["not unused by the app's rule"]
|
|
)
|
|
conn.rollback()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _off_text_node(text, save=False):
|
|
return {
|
|
"class_type": "PromptManagerText",
|
|
"inputs": {"text": text, "save_to_library": save},
|
|
}
|
|
|
|
|
|
def _encode_from(source):
|
|
return {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": source, "clip": ["1", 1]},
|
|
}
|
|
|
|
|
|
def off_behind_pass_through(branch=2):
|
|
"""The dataset graph with ``branch``'s text coming from a switched-off
|
|
PromptManagerText through a pass-through the #90 walk does not follow."""
|
|
graph = shapes.shared_prefix()
|
|
graph["50"] = _off_text_node(TEXTS[branch - 1])
|
|
graph["51"] = {"class_type": "MyNode", "inputs": {"value": ["50", 0]}}
|
|
graph[str(branch * 10)] = _encode_from(["51", 0])
|
|
return graph
|
|
|
|
|
|
def off_target(branch=2, save=False):
|
|
"""``branch``'s text straight from a PromptManagerText whose switch is
|
|
``save`` (False, or a link to a node the graph cannot read)."""
|
|
graph = shapes.shared_prefix()
|
|
graph["50"] = _off_text_node(TEXTS[branch - 1], save)
|
|
graph[str(branch * 10)] = _encode_from(["50", 0])
|
|
if isinstance(save, list):
|
|
graph["52"] = {"class_type": "SomeBoolNode", "inputs": {}}
|
|
return graph
|
|
|
|
|
|
def off_negative_only():
|
|
"""Every branch's negative from a switched-off PromptManagerText."""
|
|
graph = shapes.shared_prefix()
|
|
graph["50"] = _off_text_node(shapes.NEGATIVE)
|
|
graph["3"] = _encode_from(["50", 0])
|
|
return graph
|
|
|
|
|
|
class SaveToLibraryTests(RecordedRuns, unittest.TestCase):
|
|
"""#90: the repair never relinks a file to a switched-off node's text.
|
|
|
|
The plan decides each file through ``decide_file_prompt``; every answer
|
|
but a positive from the recorded saver or the inspector is a skip, so a
|
|
``not_saved`` or ``not_saved_wired`` file is never written. Each file is
|
|
reported by a run without contexts (as a cached or unrecorded run is),
|
|
so only the ladder can stop it. Saver "23" is branch 2.
|
|
"""
|
|
|
|
def setUp(self):
|
|
self.tmp = tempfile.mkdtemp()
|
|
self.addCleanup(shutil.rmtree, self.tmp, True)
|
|
self.addCleanup(assert_closed, self, self.tmp)
|
|
self.addCleanup(PromptModel.reset_schema_cache)
|
|
self.output = os.path.join(self.tmp, "output")
|
|
os.makedirs(os.path.join(self.output, "dataset"))
|
|
self.db_path = os.path.join(self.tmp, "prompts.db")
|
|
|
|
def plan_for(self, graph):
|
|
path = os.path.join(self.output, "dataset", NAME.format(1))
|
|
png(path, graph)
|
|
db = PromptDatabase(self.db_path)
|
|
try:
|
|
wrong = db.save_prompt(
|
|
text=TEXTS[0], prompt_hash=generate_prompt_hash(TEXTS[0])
|
|
)
|
|
db.link_image_to_prompt(wrong, path, metadata={})
|
|
self.record_run(
|
|
db, self.output, [("dataset", NAME.format(1))], [], nodes=["23"]
|
|
)
|
|
plan = tool.build_plan(db, self.output, self.db_path)
|
|
finally:
|
|
db.close_all()
|
|
(item,) = plan["rows"]
|
|
return item
|
|
|
|
def test_an_off_node_behind_a_pass_through_is_skipped(self):
|
|
item = self.plan_for(off_behind_pass_through())
|
|
self.assertEqual(item["action"], "skip")
|
|
self.assertEqual(item["decision"], "not_saved/none")
|
|
self.assertNotIn("target_text", item)
|
|
|
|
def test_an_off_node_in_the_recorded_run_target_is_skipped(self):
|
|
cases = {
|
|
"off": (False, "not_saved/none"),
|
|
"wired": (["52", 0], "not_saved_wired/none"),
|
|
}
|
|
for name, (save, decision) in cases.items():
|
|
with self.subTest(name):
|
|
if os.path.exists(self.db_path):
|
|
os.remove(self.db_path)
|
|
item = self.plan_for(off_target(save=save))
|
|
self.assertEqual(item["action"], "skip")
|
|
self.assertEqual(item["decision"], decision)
|
|
|
|
def test_an_off_node_on_another_branch_still_skips_the_file(self):
|
|
"""The ladder judges the file's whole graph: any switched-off text
|
|
upstream of a positive makes the file not saved."""
|
|
item = self.plan_for(off_target(branch=3))
|
|
self.assertEqual(item["action"], "skip")
|
|
self.assertEqual(item["decision"], "not_saved/none")
|
|
|
|
def test_an_off_node_wired_only_to_the_negative_still_repairs(self):
|
|
item = self.plan_for(off_negative_only())
|
|
self.assertEqual(item["action"], "relink")
|
|
self.assertEqual(item["source"], "recorded-run")
|
|
self.assertEqual(item["target_text"], TEXTS[1])
|
|
|
|
def test_a_switched_on_node_in_the_target_still_repairs(self):
|
|
"""Regression guard: a switch that is on is no reason to skip."""
|
|
item = self.plan_for(off_target(save=True))
|
|
self.assertEqual(item["action"], "relink")
|
|
self.assertEqual(item["target_text"], TEXTS[1])
|
|
|
|
|
|
class ActionTests(unittest.TestCase):
|
|
"""Only the recorded saver relinks; a certain step naming the linked
|
|
prompt keeps it; the legacy floor and the rest are skipped."""
|
|
|
|
def test_action_by_decision(self):
|
|
linked = generate_prompt_hash("linked")
|
|
cases = [
|
|
(Decision("positive", "other", source="recorded-run"), "relink"),
|
|
(Decision("positive", "linked", source="recorded-run"), "keep"),
|
|
(Decision("positive", "linked", source="inspector"), "keep"),
|
|
(Decision("positive", "other", source="inspector"), "skip"),
|
|
(Decision("positive", "linked", source="legacy-parser"), "skip"),
|
|
(Decision("pair", "linked", source="resolver"), "skip"),
|
|
(Decision("never"), "skip"),
|
|
(Decision("not_saved"), "skip"),
|
|
(Decision("not_saved_wired"), "skip"),
|
|
(Decision("positive", "other", source="opted-out-negative"), "skip"),
|
|
]
|
|
for decision, action in cases:
|
|
with self.subTest(decision.source, text=decision.text):
|
|
self.assertEqual(tool.action_of(decision, linked), action)
|
|
|
|
|
|
class CommandLineTests(unittest.TestCase):
|
|
def test_running_probe_reports_a_listening_port(self):
|
|
import socket
|
|
|
|
with socket.socket() as server:
|
|
server.bind(("127.0.0.1", 0))
|
|
server.listen(1)
|
|
port = server.getsockname()[1]
|
|
found = tool.comfyui_running("/nonexistent/comfy", port)
|
|
self.assertTrue(any(str(port) in item for item in found))
|
|
|
|
def test_running_probe_is_quiet_when_nothing_runs(self):
|
|
import socket
|
|
|
|
with socket.socket() as probe:
|
|
probe.bind(("127.0.0.1", 0))
|
|
port = probe.getsockname()[1]
|
|
self.assertEqual(tool.comfyui_running("/nonexistent/comfy", port), [])
|
|
|
|
def fake_proc(self, tmp):
|
|
proc = Path(tmp, "proc")
|
|
db = Path(tmp, "prompts.db")
|
|
db.write_bytes(b"")
|
|
Path(tmp, "prompts.db-wal").write_bytes(b"")
|
|
comfy = Path(tmp, "comfy")
|
|
comfy.mkdir()
|
|
holder = proc / "125" / "fd"
|
|
holder.mkdir(parents=True)
|
|
os.symlink(Path(tmp, "prompts.db-wal"), holder / "3")
|
|
(proc / "125" / "cmdline").write_bytes(b"python\0main.py\0")
|
|
os.symlink(comfy, proc / "125" / "cwd")
|
|
(proc / "123").mkdir() # unreadable: no fd, cmdline or cwd
|
|
(proc / "self").mkdir() # not a pid
|
|
return proc, db, comfy
|
|
|
|
@unittest.skipIf(os.name == "nt", "symlinks need privileges on Windows")
|
|
def test_proc_scan_finds_holders_and_skips_unreadable_entries(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
proc, db, comfy = self.fake_proc(tmp)
|
|
found = tool.open_by_others(str(db), proc=str(proc), own_pid=1)
|
|
self.assertEqual(len(found), 1)
|
|
self.assertIn("125", found[0])
|
|
self.assertEqual(
|
|
tool.open_by_others(str(db), proc=str(proc), own_pid=125), []
|
|
)
|
|
running = tool.comfyui_running(str(comfy), 1, proc=str(proc))
|
|
self.assertTrue(any("125" in item for item in running), running)
|
|
|
|
def test_without_proc_the_scan_says_it_cannot_tell(self):
|
|
self.assertEqual(
|
|
tool.open_by_others("/x/prompts.db", proc="/nonexistent/proc"),
|
|
[tool.NO_PROC],
|
|
)
|
|
|
|
def test_missing_database_is_an_argument_error(self):
|
|
with self.assertRaises(SystemExit), contextlib.redirect_stderr(io.StringIO()):
|
|
tool.main(
|
|
["/nonexistent/prompts.db", "--output-dir", "/tmp"], out=io.StringIO()
|
|
)
|
|
|
|
def test_apply_needs_the_comfyui_folder(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
db = os.path.join(tmp, "prompts.db")
|
|
Path(db).write_bytes(b"")
|
|
with (
|
|
self.assertRaises(SystemExit),
|
|
contextlib.redirect_stderr(io.StringIO()),
|
|
):
|
|
tool.main([db, "--apply", "plan.json"], out=io.StringIO())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|