Vito
e5da4e1d9d
* 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
342 lines
15 KiB
Python
342 lines
15 KiB
Python
"""Reusable pairs own their metadata without changing shared prompt content."""
|
|
|
|
import os
|
|
import sqlite3
|
|
import tempfile
|
|
import unittest
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
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.operations import PromptDatabase
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
|
|
class TestPromptPairs(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.tmp.cleanup)
|
|
self.addCleanup(assert_closed, self, self.tmp.name)
|
|
self.db = PromptDatabase(os.path.join(self.tmp.name, "pairs.db"))
|
|
self.addCleanup(self.db.close_all)
|
|
|
|
def save(self, positive="a cat", negative="blurry", **kwargs):
|
|
return self.db.save_prompt_pair(positive, negative, **kwargs)
|
|
|
|
def test_pair_dedup_is_ordered_and_empty_negative_is_explicit(self):
|
|
first = self.save()
|
|
self.assertEqual(first, self.save(" A CAT ", " BLURRY "))
|
|
self.assertNotEqual(first, self.save("blurry", "a cat"))
|
|
self.assertNotEqual(first, self.save("a cat", "watermark"))
|
|
empty = self.save("a cat", "")
|
|
self.assertEqual(empty, self.save("a cat", " "))
|
|
self.assertIsNone(self.db.get_prompt_pair(empty)["negative_prompt_id"])
|
|
self.assertEqual(self.db.get_prompt_pair(empty)["negative_text"], "")
|
|
self.assertNotEqual(empty, first)
|
|
|
|
def test_same_text_can_serve_both_roles_and_existing_metadata_survives(self):
|
|
prompt = self.db.save_prompt(
|
|
"a cat",
|
|
category="original",
|
|
tags=["old"],
|
|
prompt_hash=generate_prompt_hash("a cat"),
|
|
)
|
|
pair = self.db.get_prompt_pair(
|
|
self.save("a cat", "a cat", category="pair", tags=["new"])
|
|
)
|
|
self.assertEqual(pair["positive_prompt_id"], prompt)
|
|
self.assertEqual(pair["negative_prompt_id"], prompt)
|
|
self.assertEqual(pair["tags"], ["new"])
|
|
self.assertEqual(self.db.get_prompt_by_id(prompt)["category"], "original")
|
|
self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["old"])
|
|
|
|
def test_repeated_save_preserves_pair_metadata_and_explicit_update_changes_it(self):
|
|
pid = self.save(
|
|
title="Title",
|
|
category="Portrait",
|
|
tags=["studio", "studio"],
|
|
rating=4,
|
|
notes="keep",
|
|
)
|
|
self.assertEqual(pid, self.save(title="overwrite", tags=["bad"], rating=1))
|
|
pair = self.db.get_prompt_pair(pid)
|
|
self.assertEqual(
|
|
(
|
|
pair["title"],
|
|
pair["category"],
|
|
pair["tags"],
|
|
pair["rating"],
|
|
pair["notes"],
|
|
),
|
|
("Title", "Portrait", ["studio"], 4, "keep"),
|
|
)
|
|
self.assertTrue(
|
|
self.db.update_prompt_pair_metadata(
|
|
pid, title="", category="", tags=[], notes=""
|
|
)
|
|
)
|
|
pair = self.db.get_prompt_pair(pid)
|
|
self.assertEqual(
|
|
(pair["title"], pair["category"], pair["tags"], pair["notes"]),
|
|
("", "", [], ""),
|
|
)
|
|
self.assertEqual(pair["rating"], 4)
|
|
self.assertFalse(self.db.update_prompt_pair_metadata(999999, title="missing"))
|
|
self.assertFalse(self.db.update_prompt_pair_metadata(pid))
|
|
|
|
def test_list_filters_pagination_archive_and_pair_usage_are_independent(self):
|
|
first = self.save(category="Animals", tags=["studio", "cat"])
|
|
second = self.save("a dog", "blur", category="Animals", tags=["studio"])
|
|
other = self.save("sky", "", category="Outdoor")
|
|
result = self.db.list_prompt_pairs(category="Animals", tags=["studio"], limit=1)
|
|
self.assertEqual(result["total"], 2)
|
|
self.assertTrue(result["has_more"])
|
|
self.assertEqual(len(result["pairs"]), 1)
|
|
self.assertEqual(
|
|
self.db.list_prompt_pairs(text="dog")["pairs"][0]["id"], second
|
|
)
|
|
self.assertEqual(self.db.list_prompt_pairs(tags=["studio", "cat"])["total"], 1)
|
|
self.assertTrue(self.db.archive_prompt_pair(first))
|
|
self.assertEqual(self.db.list_prompt_pairs()["total"], 2)
|
|
self.assertEqual(self.db.list_prompt_pairs(include_archived=True)["total"], 3)
|
|
self.assertTrue(self.db.get_prompt_pair(first)["archived"])
|
|
self.assertTrue(self.db.archive_prompt_pair(first, False))
|
|
self.assertTrue(self.db.record_prompt_pair_use(other, 2))
|
|
pair = self.db.get_prompt_pair(other)
|
|
self.assertEqual(pair["run_count"], 2)
|
|
self.assertIsNotNone(pair["last_used_at"])
|
|
self.assertEqual(
|
|
self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 0
|
|
)
|
|
|
|
def test_failure_during_tag_save_rolls_back_pair_and_components(self):
|
|
with patch.object(
|
|
self.db, "_ensure_tags", side_effect=sqlite3.OperationalError("injected")
|
|
):
|
|
with self.assertRaises(sqlite3.OperationalError):
|
|
self.save(tags=["studio"])
|
|
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|
|
self.assertEqual(self.db.search_prompts(), [])
|
|
|
|
def test_concurrent_same_pair_returns_one_id_and_preserves_first_metadata(self):
|
|
first = self.save(title="original")
|
|
|
|
def save(index):
|
|
try:
|
|
return self.save(title=f"writer {index}")
|
|
finally:
|
|
self.db.close()
|
|
|
|
with ThreadPoolExecutor(max_workers=6) as executor:
|
|
ids = list(executor.map(save, range(24)))
|
|
self.assertEqual(set(ids), {first})
|
|
self.assertEqual(self.db.get_prompt_pair(first)["title"], "original")
|
|
self.assertEqual(self.db.list_prompt_pairs()["total"], 1)
|
|
|
|
def test_concurrent_empty_negative_pair_creation_is_atomic(self):
|
|
def save(_):
|
|
try:
|
|
return self.save(negative="")
|
|
finally:
|
|
self.db.close()
|
|
|
|
with ThreadPoolExecutor(max_workers=6) as executor:
|
|
ids = list(executor.map(save, range(12)))
|
|
self.assertEqual(len(set(ids)), 1)
|
|
self.assertEqual(len(self.db.search_prompts()), 1)
|
|
|
|
def test_component_edit_and_all_delete_paths_reject_dependencies_atomically(self):
|
|
from database.operations import PairDependencyError
|
|
|
|
pair = self.db.get_prompt_pair(self.save())
|
|
positive = pair["positive_prompt_id"]
|
|
image = self.db.link_image_to_prompt(
|
|
positive, os.path.join(self.tmp.name, "keep.png")
|
|
)
|
|
unused = self.db.save_prompt(
|
|
"unused", category="remove", prompt_hash=generate_prompt_hash("unused")
|
|
)
|
|
self.db.update_prompt_metadata(positive, category="remove")
|
|
for action in (
|
|
lambda: self.db.update_prompt_text(positive, "replacement"),
|
|
lambda: self.db.delete_prompt(positive),
|
|
lambda: self.db.bulk_delete_prompts([unused, positive]),
|
|
lambda: self.db.delete_prompts_by_category("remove"),
|
|
):
|
|
with self.subTest(action=action):
|
|
with self.assertRaises(PairDependencyError):
|
|
action()
|
|
self.assertIsNotNone(self.db.get_prompt_by_id(unused))
|
|
self.assertEqual(self.db.get_prompt_by_id(positive)["text"], "a cat")
|
|
self.assertIsNotNone(self.db.get_image_by_id(image))
|
|
self.assertTrue(self.db.delete_prompt_pair(pair["id"]))
|
|
self.assertTrue(self.db.update_prompt_text(positive, "replacement"))
|
|
|
|
def test_prune_skips_members_even_in_archived_pairs_and_preserves_other_behavior(
|
|
self,
|
|
):
|
|
pair = self.db.get_prompt_pair(self.save())
|
|
self.db.archive_prompt_pair(pair["id"])
|
|
unused = self.db.save_prompt(
|
|
"unused", prompt_hash=generate_prompt_hash("unused")
|
|
)
|
|
self.assertEqual(self.db.prune_orphaned_prompts(), 1)
|
|
self.assertIsNone(self.db.get_prompt_by_id(unused))
|
|
self.assertIsNotNone(self.db.get_prompt_by_id(pair["negative_prompt_id"]))
|
|
|
|
def test_duplicate_cleanup_skips_referenced_group_without_losing_pairs(self):
|
|
pair = self.db.get_prompt_pair(self.save())
|
|
with self.db.model.get_connection() as conn:
|
|
conn.execute(
|
|
"INSERT INTO prompts(text,category) VALUES(' A CAT ', 'duplicate')"
|
|
)
|
|
self.assertEqual(self.db.cleanup_duplicates(), 0)
|
|
self.assertEqual(self.db.get_prompt_pair(pair["id"])["positive_text"], "a cat")
|
|
self.assertEqual(len(self.db.search_prompts()), 3)
|
|
|
|
def test_image_association_preserves_provenance_and_pair_deletion_requires_archive(
|
|
self,
|
|
):
|
|
from database.operations import PairDependencyError
|
|
|
|
pid = self.save()
|
|
pair = self.db.get_prompt_pair(pid)
|
|
image = self.db.link_image_to_prompt(
|
|
pair["positive_prompt_id"], os.path.join(self.tmp.name, "render.png")
|
|
)
|
|
provenance = {
|
|
"run_id": "abc",
|
|
"negative_text": "actual expanded negative",
|
|
"positive_text": "trigger a cat",
|
|
}
|
|
association = self.db.link_image_to_pair(pid, image, provenance)
|
|
self.assertEqual(
|
|
association,
|
|
self.db.link_image_to_pair(
|
|
pid, image, dict(reversed(list(provenance.items())))
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
self.db.get_image_prompt_pairs(image)[0]["provenance"], provenance
|
|
)
|
|
self.assertEqual(self.db.get_prompt_pair(pid)["image_count"], 1)
|
|
with self.assertRaises(PairDependencyError):
|
|
self.db.delete_prompt_pair(pid)
|
|
self.assertTrue(self.db.archive_prompt_pair(pid))
|
|
with self.assertRaises(ValueError):
|
|
self.db.link_image_to_pair(99999, image, provenance)
|
|
self.assertTrue(self.db.delete_image(image))
|
|
self.assertEqual(self.db.get_image_prompt_pairs(image), [])
|
|
self.assertTrue(self.db.delete_prompt_pair(pid))
|
|
self.assertIsNotNone(self.db.get_prompt_by_id(pair["positive_prompt_id"]))
|
|
|
|
def test_pair_validates_values_before_writing(self):
|
|
for positive, negative, kwargs in (
|
|
("", "blur", {}),
|
|
("cat", None, {}),
|
|
("cat", "blur", {"rating": 6}),
|
|
("cat", "blur", {"tags": "one"}),
|
|
):
|
|
with self.subTest(positive=positive, negative=negative, kwargs=kwargs):
|
|
with self.assertRaises((TypeError, ValueError)):
|
|
self.save(positive, negative, **kwargs)
|
|
self.assertEqual(self.db.search_prompts(), [])
|
|
for kwargs in ({"limit": 0}, {"offset": -1}, {"limit": True}):
|
|
with self.assertRaises(ValueError):
|
|
self.db.list_prompt_pairs(**kwargs)
|
|
self.assertIsNone(self.db.get_prompt_pair(99999))
|
|
self.assertFalse(self.db.delete_prompt_pair(99999))
|
|
|
|
def test_connection_helpers_leave_commit_and_rollback_to_the_caller(self):
|
|
conn = self.db.model.get_connection()
|
|
|
|
def rollback_transaction():
|
|
with conn:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
self.db._save_pair(conn, "transient", "temporary", tags=["tag"])
|
|
self.assertRaisesRegex(RuntimeError, "own transaction", self.save)
|
|
raise RuntimeError("rollback")
|
|
|
|
self.assertRaisesRegex(RuntimeError, "rollback", rollback_transaction)
|
|
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|
|
self.assertEqual(self.db.search_prompts(), [])
|
|
|
|
def test_metadata_update_failure_rolls_back_prior_fields_and_tags(self):
|
|
pid = self.save(title="keep", tags=["old"])
|
|
with patch.object(
|
|
self.db, "_ensure_tags", side_effect=sqlite3.OperationalError("injected")
|
|
):
|
|
with self.assertRaises(sqlite3.OperationalError):
|
|
self.db.update_prompt_pair_metadata(pid, title="replace", tags=["new"])
|
|
pair = self.db.get_prompt_pair(pid)
|
|
self.assertEqual(pair["title"], "keep")
|
|
self.assertEqual(pair["tags"], ["old"])
|
|
|
|
def test_invalid_associations_metadata_and_missing_mutations_are_safe(self):
|
|
pid = self.save()
|
|
other = self.db.save_prompt("other", prompt_hash=generate_prompt_hash("other"))
|
|
image = self.db.link_image_to_prompt(
|
|
other, os.path.join(self.tmp.name, "other.png")
|
|
)
|
|
for provenance in ([], {"number": float("nan")}, {"object": object()}):
|
|
with self.assertRaises(ValueError):
|
|
self.db.link_image_to_pair(pid, image, provenance)
|
|
with self.assertRaisesRegex(ValueError, "positive prompt"):
|
|
self.db.link_image_to_pair(pid, image, {})
|
|
self.assertEqual(self.db.get_image_prompt_pairs(image), [])
|
|
self.assertFalse(self.db.archive_prompt_pair(999999))
|
|
self.assertFalse(self.db.record_prompt_pair_use(999999))
|
|
for times in (0, True):
|
|
with self.assertRaises(ValueError):
|
|
self.db.record_prompt_pair_use(pid, times)
|
|
with self.assertRaises(ValueError):
|
|
self.db.archive_prompt_pair(pid, "yes")
|
|
with self.assertRaises(ValueError):
|
|
self.db.update_prompt_pair_metadata(pid, title=[])
|
|
with self.assertRaises(ValueError):
|
|
self.db.list_prompt_pairs(text=[])
|
|
|
|
def test_connection_outcomes_report_insertions_under_the_held_transaction(self):
|
|
conn = self.db.model.get_connection()
|
|
legacy = self.save("Kept Spelling", "", title="legacy", tags=["t"])
|
|
with conn:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
first = self.db._pair_component_outcome(conn, " New Text ")
|
|
again = self.db._pair_component_outcome(conn, "new text")
|
|
spaced = self.db._pair_component_outcome(conn, "new text")
|
|
kept = self.db._pair_component_outcome(conn, " kept spelling ")
|
|
pair = self.db._save_pair_components(conn, first[0], spaced[0], title="p")
|
|
reused = self.db._save_pair_components(conn, first[0], spaced[0])
|
|
empty = self.db._save_pair_components(conn, kept[0], None, title="x")
|
|
self.assertEqual((first[1], again, spaced[1]), (True, (first[0], False), True))
|
|
self.assertNotEqual(first[0], spaced[0])
|
|
self.assertEqual(kept[1], False)
|
|
self.assertEqual((pair[1], reused), (True, (pair[0], False)))
|
|
self.assertEqual(empty, (legacy, False))
|
|
self.assertEqual(self.db.get_prompt_by_id(first[0])["text"], "New Text")
|
|
self.assertEqual(self.db.get_prompt_pair(pair[0])["title"], "p")
|
|
legacy_pair = self.db.get_prompt_pair(legacy)
|
|
self.assertEqual((legacy_pair["title"], legacy_pair["tags"]), ("legacy", ["t"]))
|
|
self.assertEqual(legacy_pair["positive_text"], "Kept Spelling")
|
|
# Public saves keep their integer contract on create and on reuse.
|
|
self.assertIs(type(self.save("fresh", "")), int)
|
|
self.assertEqual(self.save("FRESH ", " "), self.save("fresh", ""))
|
|
|
|
def test_pair_component_outcome_rolls_back_with_its_caller(self):
|
|
conn = self.db.model.get_connection()
|
|
|
|
def failing():
|
|
with conn:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
positive, _ = self.db._pair_component_outcome(conn, "transient")
|
|
self.db._save_pair_components(conn, positive, None, tags=["x"])
|
|
raise RuntimeError("rollback")
|
|
|
|
self.assertRaisesRegex(RuntimeError, "rollback", failing)
|
|
self.assertEqual(self.db.search_prompts(), [])
|
|
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|