Files
ComfyAssets-ComfyUI_PromptM…/tests/test_pair_schema_restore.py
Vito e5da4e1d9d feat(library): backend for native recipe editing, pair archive browsing and portable export (#233)
* feat(library): persist recoverable parent state

* feat(library): expose versioned parent metadata reads

* feat(library): protect parent metadata and archive writes

* feat(library): save native recipes with atomic outcomes

* feat(library): download complete portable library json

* fix(library): name export owners list positives per spec

* fix(library): normalise native category and harden restore and id errors

* fix(library): measure stored titles past embedded NULs
2026-10-18 20:03:15 -07:00

330 lines
14 KiB
Python

"""Pair schema installation and restore failures use disposable databases only."""
import os
import sqlite3
import tempfile
import unittest
from contextlib import closing
from pathlib import Path
from unittest.mock import patch
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
class PairSchemaTests(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.addCleanup(self.temp.cleanup)
self.addCleanup(assert_closed, self, self.temp.name)
self.addCleanup(PromptModel.reset_schema_cache)
self.models = []
self.addCleanup(self.close_models)
self.path = str(Path(self.temp.name) / "live.db")
def close_models(self):
for model in self.models:
model.close_all()
def model(self, path=None):
model = PromptModel(path or self.path)
self.models.append(model)
return model
def populate(self, model):
with model.get_connection() as conn:
conn.execute(
"INSERT INTO prompts(id,text,hash,category,notes,run_count) VALUES"
"(7,'positive','p','keep','notes',12)"
)
conn.execute("INSERT INTO prompts(id,text,hash) VALUES(9,'negative','n')")
conn.execute(
"INSERT INTO generated_images(id,prompt_id,image_path,filename,fil"
"e_path) VALUES(4,7,'/tmp/image.png','image.png','/tmp/image.png')"
)
conn.execute("INSERT INTO tags(id,name) VALUES(3,'keep-tag')")
conn.execute("INSERT INTO prompt_tags VALUES(7,3)")
def legacy(self):
model = self.model(str(Path(self.temp.name) / "legacy.db"))
self.populate(model)
with model.get_connection() as conn:
for table in ("image_prompt_pairs", "prompt_pair_tags", "prompt_pairs"):
conn.execute(f"DROP TABLE IF EXISTS {table}")
return model
def test_restore_scenarios_close_all_sqlite_handles(self):
connections = []
connect = sqlite3.connect
def tracked_connect(*args, **kwargs):
connection = connect(*args, **kwargs)
connections.append(connection)
return connection
scenarios = (
"test_pair_ddl_failure_rolls_back_and_retry_succeeds",
"test_restore_legacy_wal_migrates_stage_and_preserves_safety_backup",
"test_restore_migration_failure_keeps_original_and_source_untouched",
"test_verification_rejects_partial_pair_schema_and_orphan_pair",
)
try:
with patch.object(sqlite3, "connect", side_effect=tracked_connect):
result = unittest.TestResult()
unittest.TestSuite(type(self)(name) for name in scenarios).run(result)
self.assertTrue(result.wasSuccessful(), result.errors + result.failures)
for index, connection in enumerate(connections):
with self.subTest(connection=index):
with self.assertRaises(sqlite3.ProgrammingError):
connection.execute("SELECT 1")
finally:
for connection in connections:
connection.close()
def test_fresh_schema_constraints_and_foreign_keys(self):
model = self.model()
self.populate(model)
with model.get_connection() as conn:
conn.execute(
"INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id"
") VALUES(2,7,9)"
)
conn.execute("INSERT INTO prompt_pairs(positive_prompt_id) VALUES(7)")
conn.execute("INSERT INTO prompt_pair_tags VALUES(2,3)")
conn.execute(
"INSERT INTO image_prompt_pairs(pair_id,image_id,provenance) VALUE"
"S(2,4,'{}')"
)
for sql in (
"INSERT INTO prompt_pairs(positive_prompt_id,negative_prompt_id) V"
"ALUES(7,9)",
"INSERT INTO prompt_pairs(positive_prompt_id) VALUES(7)",
"INSERT INTO prompt_pairs(positive_prompt_id) VALUES(404)",
"INSERT INTO prompt_pairs(positive_prompt_id,rating) VALUES(9,6)",
"INSERT INTO prompt_pairs(positive_prompt_id,archived) VALUES(9,2)",
"DELETE FROM prompts WHERE id=9",
"DELETE FROM prompt_pairs WHERE id=2",
"INSERT INTO image_prompt_pairs(pair_id,image_id,provenance) VALUE"
"S(2,4,'{}')",
):
with self.subTest(sql=sql), self.assertRaises(sqlite3.IntegrityError):
with model.get_connection() as conn:
conn.execute(sql)
with model.get_connection() as conn:
conn.execute(
"INSERT INTO prompt_pairs(positive_prompt_id,negative_prompt_id) V"
"ALUES(9,7)"
)
conn.execute("DELETE FROM generated_images WHERE id=4")
self.assertEqual(
conn.execute("SELECT COUNT(*) FROM image_prompt_pairs").fetchone()[0], 0
)
self.assertEqual(conn.execute("PRAGMA foreign_key_check").fetchall(), [])
def test_legacy_upgrade_preserves_ids_metadata_images_counts_and_repeats(self):
legacy = self.legacy()
path = legacy.db_path
legacy.close_all()
PromptModel.reset_schema_cache()
for _ in range(2):
model = self.model(path)
with model.get_connection() as conn:
self.assertEqual(
tuple(
conn.execute(
"SELECT id,text,category,notes,run_count FROM prompts WHER"
"E id=7"
).fetchone()
),
(7, "positive", "keep", "notes", 12),
)
self.assertEqual(
conn.execute("SELECT id FROM generated_images").fetchone()[0], 4
)
self.assertEqual(
tuple(conn.execute("SELECT * FROM prompt_tags").fetchone()), (7, 3)
)
self.assertEqual(
conn.execute("SELECT COUNT(*) FROM prompt_pairs").fetchone()[0], 0
)
model.close_all()
PromptModel.reset_schema_cache()
def test_pair_ddl_failure_rolls_back_and_retry_succeeds(self):
with patch.object(
PromptModel,
"_create_pair_indexes",
side_effect=sqlite3.OperationalError("injected index failure"),
):
with self.assertRaisesRegex(sqlite3.OperationalError, "injected"):
self.model()
with closing(sqlite3.connect(self.path)) as conn, conn:
self.assertEqual(
conn.execute(
"SELECT name FROM sqlite_master WHERE name IN ('prompt_pairs','pro"
"mpt_pair_tags','image_prompt_pairs')"
).fetchall(),
[],
)
model = self.model()
self.assertTrue(PromptModel.verify_database_file(model.db_path)[0])
def test_verification_rejects_partial_pair_schema_and_orphan_pair(self):
model = self.model()
self.populate(model)
with model.get_connection() as conn:
conn.execute(
"INSERT INTO prompt_pairs(positive_prompt_id,negative_prompt_id) V"
"ALUES(7,9)"
)
model.close_all()
with closing(sqlite3.connect(model.db_path)) as conn, conn:
conn.execute("PRAGMA foreign_keys=OFF")
conn.execute("UPDATE prompt_pairs SET negative_prompt_id=404")
ok, reason = PromptModel.verify_database_file(model.db_path)
self.assertFalse(ok)
self.assertIn("foreign key", reason.lower())
with closing(sqlite3.connect(model.db_path)) as conn, conn:
conn.execute("DROP TABLE prompt_pair_tags")
ok, reason = PromptModel.verify_database_file(model.db_path)
self.assertFalse(ok)
self.assertIn("pair", reason.lower())
def test_restore_migration_failure_keeps_original_and_source_untouched(self):
live = self.model()
self.populate(live)
legacy = self.legacy()
with live.get_connection() as conn:
conn.execute("UPDATE prompts SET text='original live text' WHERE id=7")
with patch.object(
PromptModel,
"_migrate_prompt_pairs",
side_effect=sqlite3.OperationalError("migration failed"),
):
with self.assertRaisesRegex(sqlite3.OperationalError, "migration failed"):
live.restore_from_file(legacy.db_path)
with live.get_connection() as conn:
self.assertEqual(
conn.execute("SELECT text FROM prompts WHERE id=7").fetchone()[0],
"original live text",
)
self.assertEqual(
conn.execute("SELECT run_count FROM prompts WHERE id=7").fetchone()[0],
12,
)
self.assertEqual(
conn.execute("SELECT COUNT(*) FROM generated_images").fetchone()[0], 1
)
with closing(sqlite3.connect(legacy.db_path)) as conn, conn:
self.assertIsNone(
conn.execute(
"SELECT name FROM sqlite_master WHERE name='prompt_pairs'"
).fetchone()
)
def test_restore_legacy_wal_migrates_stage_and_preserves_safety_backup(self):
live = self.model()
with live.get_connection() as conn:
conn.execute("INSERT INTO prompts(text,hash) VALUES('original','original')")
legacy = self.legacy()
with legacy.get_connection() as conn:
conn.execute("INSERT INTO prompts(text,hash) VALUES('wal-only','wal-only')")
self.assertGreater(os.path.getsize(legacy.db_path + "-wal"), 0)
backup = live.restore_from_file(legacy.db_path)
self.assertTrue(Path(backup).is_file())
with closing(sqlite3.connect(backup)) as conn, conn:
self.assertEqual(
conn.execute("SELECT text FROM prompts").fetchall(), [("original",)]
)
with live.get_connection() as conn:
self.assertEqual(
conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0], 3
)
self.assertEqual(
conn.execute("SELECT COUNT(*) FROM prompt_pairs").fetchone()[0], 0
)
self.assertEqual(
[r[1] for r in conn.execute("PRAGMA table_info(prompt_library_state)")],
["prompt_id", "title"],
)
self.assertTrue(PromptModel.verify_database_file(live.db_path)[0])
def test_restore_staged_verification_failure_does_not_replace_live_pages(self):
live = self.model()
self.populate(live)
source = self.legacy()
with patch.object(
PromptModel,
"verify_database_file",
side_effect=[(True, "ok"), (False, "staged validation failed")],
):
with self.assertRaisesRegex(ValueError, "staged validation failed"):
live.restore_from_file(source.db_path)
with live.get_connection() as conn:
self.assertEqual(
conn.execute("SELECT run_count FROM prompts WHERE id=7").fetchone()[0],
12,
)
def test_verification_rejects_missing_null_uniqueness(self):
model = self.model()
with model.get_connection() as conn:
conn.execute("DROP INDEX idx_prompt_pairs_positive_only")
ok, reason = PromptModel.verify_database_file(model.db_path)
self.assertFalse(ok)
self.assertIn("uniqueness", reason)
def test_restore_round_trips_pair_metadata_tags_and_provenance(self):
source = self.model(str(Path(self.temp.name) / "paired.db"))
self.populate(source)
with source.get_connection() as conn:
conn.execute(
"INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id"
",title,notes,run_count,archived) VALUES(6,7,9,'pair title','pair "
"notes',19,1)"
)
conn.execute("INSERT INTO prompt_pair_tags VALUES(6,3)")
conn.execute(
"INSERT INTO image_prompt_pairs(id,pair_id,image_id,provenance) VA"
'LUES(8,6,4,\'{"execution":"kept"}\')'
)
live = self.model()
live.restore_from_file(source.db_path)
with live.get_connection() as conn:
self.assertEqual(
tuple(
conn.execute(
"SELECT id,title,notes,run_count,archived FROM prompt_pairs"
).fetchone()
),
(6, "pair title", "pair notes", 19, 1),
)
self.assertEqual(
tuple(conn.execute("SELECT * FROM prompt_pair_tags").fetchone()), (6, 3)
)
self.assertEqual(
conn.execute(
"SELECT provenance FROM image_prompt_pairs WHERE id=8"
).fetchone()[0],
'{"execution":"kept"}',
)
def test_restore_rejects_trigger_before_running_migrations(self):
live = self.model()
source = self.legacy()
with source.get_connection() as conn:
conn.execute(
"CREATE TRIGGER dangerous AFTER INSERT ON prompts BEGIN DELETE FRO"
"M prompts; END"
)
with self.assertRaisesRegex(ValueError, "trigger"):
live.restore_from_file(source.db_path)
if __name__ == "__main__":
unittest.main()