Vito
fdba6fd012
* feat(workspace): group saved negative variants in native library (#208) * fix(library): preserve full name order and batch previews (#208)
300 lines
13 KiB
Python
300 lines
13 KiB
Python
"""Grouped library reads retain component roles and association semantics."""
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from database.operations import PromptDatabase
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
|
|
class TestLibrary(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.db = PromptDatabase(str(Path(self.temp.name) / "library.db"))
|
|
|
|
def tearDown(self):
|
|
self.db.close_all()
|
|
self.temp.cleanup()
|
|
|
|
def save(self, text, **metadata):
|
|
return self.db.save_prompt(
|
|
text, prompt_hash=generate_prompt_hash(text), **metadata
|
|
)
|
|
|
|
def test_groups_positive_parents_excludes_negative_only_and_keeps_dual_once(self):
|
|
first = self.db.save_prompt_pair("Cat", "Blur")
|
|
self.db.save_prompt_pair("Cat", "Noise")
|
|
self.db.save_prompt_pair("Blur", "Watermark")
|
|
legacy = self.save("Legacy")
|
|
rows = self.db.list_library()["items"]
|
|
self.assertEqual({row["text"] for row in rows}, {"Cat", "Blur", "Legacy"})
|
|
parent = next(row for row in rows if row["text"] == "Cat")
|
|
self.assertEqual(parent["variant_count"], 2)
|
|
self.assertEqual(len([row for row in rows if row["id"] == legacy]), 1)
|
|
self.db.archive_prompt_pair(first)
|
|
detail = self.db.get_library_positive(parent["id"])
|
|
self.assertEqual(len(detail["variants"]), 1)
|
|
self.assertEqual(detail["positive"]["variant_count"], 1)
|
|
|
|
def test_filters_match_all_constraints_on_same_parent_or_same_active_pair(self):
|
|
a = self.db.save_prompt_pair("Cat", "Blur", category="A", tags=["red"])
|
|
self.db.save_prompt_pair("Cat", "Noise", category="B", tags=["blue"])
|
|
self.assertEqual(self.db.list_library(category="A", tags=["blue"])["items"], [])
|
|
self.assertEqual(
|
|
self.db.list_library(text="Blur", category="A", tags=["red"])["total"], 1
|
|
)
|
|
self.db.archive_prompt_pair(a)
|
|
self.assertEqual(self.db.list_library(text="Blur")["total"], 0)
|
|
self.save("Legacy cat", category="A", tags=["red", "blue"])
|
|
self.assertEqual(
|
|
self.db.list_library(category="A", tags=["red", "blue"])["total"], 1
|
|
)
|
|
self.assertEqual(self.db.list_library(text="%")["items"], [])
|
|
|
|
def test_keyset_handles_insertion_deletion_ties_and_invalid_parameters(self):
|
|
for text in ("a", "b", "c", "d"):
|
|
self.save(text)
|
|
first = self.db.list_library(limit=2)
|
|
self.save("new")
|
|
second = self.db.list_library(limit=2, after=first["next_key"])
|
|
self.assertEqual([row["text"] for row in second["items"]], ["b", "a"])
|
|
self.assertFalse(second["has_more"])
|
|
for options in ({"limit": 0}, {"sort": "sql"}, {"tags": [2]}, {"text": 2}):
|
|
with self.assertRaises(ValueError):
|
|
self.db.list_library(**options)
|
|
|
|
def test_detail_distinguishes_explicit_empty_and_deduplicates_provenance_rows(self):
|
|
pair = self.db.save_prompt_pair("Cat", "")
|
|
positive = self.db.get_prompt_pair(pair)["positive_prompt_id"]
|
|
image = self.db.link_image_to_prompt(
|
|
positive, str(Path(self.temp.name) / "image.png")
|
|
)
|
|
self.db.link_image_to_pair(pair, image, {"source": "one"})
|
|
self.db.link_image_to_pair(pair, image, {"source": "two"})
|
|
self.db.link_image_to_prompt(
|
|
positive, str(Path(self.temp.name) / "unassigned.png")
|
|
)
|
|
detail = self.db.get_library_positive(positive)
|
|
self.assertEqual(detail["positive"]["image_count"], 2)
|
|
self.assertEqual(detail["unassigned_image_count"], 1)
|
|
self.assertEqual(detail["variants"][0]["negative_text"], "")
|
|
self.assertTrue(detail["variants"][0]["explicit_empty"])
|
|
self.assertEqual(detail["variants"][0]["image_count"], 1)
|
|
page = self.db.list_library_images(positive, pair_id=pair)
|
|
self.assertEqual(len(page["images"]), 1)
|
|
self.assertEqual(page["total"], 1)
|
|
unrelated = self.db.save_prompt_pair("Dog", "")
|
|
with self.assertRaises(ValueError):
|
|
self.db.list_library_images(positive, pair_id=unrelated)
|
|
self.assertIsNone(self.db.get_library_positive(9999))
|
|
|
|
def test_options_include_active_pair_categories_not_negative_only_or_archived(self):
|
|
self.save("Legacy", category="Legacy category")
|
|
pair = self.db.save_prompt_pair("Cat", "Noise", category="Pair category")
|
|
negative = self.db.get_prompt_pair(pair)["negative_prompt_id"]
|
|
with self.db.model.get_connection() as connection:
|
|
connection.execute(
|
|
"UPDATE prompts SET category='Negative only' WHERE id=?", (negative,)
|
|
)
|
|
archived = self.db.save_prompt_pair("Dog", "Blur", category="Archived category")
|
|
self.db.archive_prompt_pair(archived)
|
|
self.assertEqual(
|
|
self.db.get_library_options(),
|
|
{"categories": ["Legacy category", "Pair category"]},
|
|
)
|
|
|
|
def test_variant_detail_is_bounded_and_cursor_continues_without_duplicates(self):
|
|
for index in range(85):
|
|
self.db.save_prompt_pair("Cat", f"Negative {index}", tags=[f"tag{index}"])
|
|
positive = self.db.list_library()["items"][0]["id"]
|
|
first = self.db.get_library_positive(positive, variant_limit=40)
|
|
self.assertEqual(len(first["variants"]), 40)
|
|
self.assertEqual(first["variant_total"], 85)
|
|
self.assertTrue(first["variants_has_more"])
|
|
second = self.db.get_library_positive(
|
|
positive, variant_limit=40, variant_after=first["variants_next_key"]
|
|
)
|
|
third = self.db.get_library_positive(
|
|
positive, variant_after=second["variants_next_key"]
|
|
)
|
|
all_variants = first["variants"] + second["variants"] + third["variants"]
|
|
self.assertEqual(len({item["id"] for item in all_variants}), 85)
|
|
self.assertFalse(third["variants_has_more"])
|
|
self.assertTrue(all(len(item["tags"]) == 1 for item in all_variants))
|
|
with self.assertRaises(ValueError):
|
|
self.db.get_library_positive(positive, variant_limit=201)
|
|
|
|
def test_sorts_keysets_survive_anchor_deletion_and_insertions_without_writes(self):
|
|
import sqlite3
|
|
|
|
a = self.save("Alpha", rating=2, tags=["tag"])
|
|
b = self.save("beta", rating=5)
|
|
c = self.save("charlie")
|
|
for sort, expected in (
|
|
("newest", [c, b, a]),
|
|
("oldest", [a, b, c]),
|
|
("name", [a, b, c]),
|
|
("rating", [b, a, c]),
|
|
):
|
|
with self.subTest(sort=sort):
|
|
first = self.db.list_library(sort=sort, limit=1)
|
|
second = self.db.list_library(
|
|
sort=sort, limit=1, after=first["next_key"]
|
|
)
|
|
third = self.db.list_library(
|
|
sort=sort, limit=1, after=second["next_key"]
|
|
)
|
|
self.assertEqual(
|
|
[part["items"][0]["id"] for part in (first, second, third)],
|
|
expected,
|
|
)
|
|
first = self.db.list_library(sort="name", limit=1)
|
|
with self.db.model.get_connection() as connection:
|
|
connection.execute("DELETE FROM prompts WHERE id=?", (a,))
|
|
self.save("between")
|
|
second = self.db.list_library(sort="name", after=first["next_key"])
|
|
self.assertEqual([item["id"] for item in second["items"]], [b, c])
|
|
connection = self.db.model.get_connection()
|
|
before = connection.total_changes
|
|
denied = {
|
|
sqlite3.SQLITE_INSERT,
|
|
sqlite3.SQLITE_UPDATE,
|
|
sqlite3.SQLITE_DELETE,
|
|
sqlite3.SQLITE_CREATE_TABLE,
|
|
sqlite3.SQLITE_ALTER_TABLE,
|
|
}
|
|
connection.set_authorizer(
|
|
lambda operation, *_: (
|
|
sqlite3.SQLITE_DENY if operation in denied else sqlite3.SQLITE_OK
|
|
)
|
|
)
|
|
try:
|
|
self.db.list_library()
|
|
self.db.get_library_positive(b)
|
|
self.db.get_library_options()
|
|
self.db.get_library_previews([b])
|
|
self.db.list_library_images(b)
|
|
finally:
|
|
connection.set_authorizer(None)
|
|
self.assertEqual(connection.total_changes, before)
|
|
|
|
def test_archived_associations_and_legacy_positive_roles_remain_distinct(self):
|
|
pair = self.db.save_prompt_pair("Cat", "Noise")
|
|
pair_data = self.db.get_prompt_pair(pair)
|
|
negative = pair_data["negative_prompt_id"]
|
|
self.db.link_image_to_prompt(negative, str(Path(self.temp.name) / "legacy.png"))
|
|
rows = self.db.list_library()["items"]
|
|
self.assertIn(negative, [row["id"] for row in rows])
|
|
positive = pair_data["positive_prompt_id"]
|
|
image = self.db.link_image_to_prompt(
|
|
positive, str(Path(self.temp.name) / "pair.png")
|
|
)
|
|
self.db.link_image_to_pair(pair, image, {"source": "test"})
|
|
self.db.archive_prompt_pair(pair)
|
|
detail = self.db.get_library_positive(positive)
|
|
self.assertEqual(detail["variants"], [])
|
|
self.assertEqual(detail["positive"]["image_count"], 1)
|
|
self.assertEqual(detail["unassigned_image_count"], 0)
|
|
self.assertEqual(self.db.list_library_images(positive)["total"], 1)
|
|
with self.assertRaises(ValueError):
|
|
self.db.list_library_images(positive, pair_id=pair)
|
|
|
|
def test_invalid_database_inputs_fail_before_queries(self):
|
|
for options in (
|
|
{"category": 2},
|
|
{"tags": [""]},
|
|
{"text": "x" * 501},
|
|
{"after": {}},
|
|
{"after": {"ceiling": "x", "id": 1, "value": 1}},
|
|
):
|
|
with self.assertRaises(ValueError):
|
|
self.db.list_library(**options)
|
|
for operation in (
|
|
lambda: self.db.get_library_positive(1, variant_after=0),
|
|
lambda: self.db.list_library_images(1, after=0),
|
|
):
|
|
with self.assertRaises(ValueError):
|
|
operation()
|
|
|
|
def test_cursor_key_comes_from_read_page_even_if_anchor_is_deleted_during_read(
|
|
self,
|
|
):
|
|
from unittest.mock import patch
|
|
|
|
self.save("a")
|
|
self.save("b")
|
|
original = self.db._library_parents
|
|
|
|
def remove_anchor(*args, **kwargs):
|
|
rows = original(*args, **kwargs)
|
|
with self.db.model.get_connection() as connection:
|
|
connection.execute("DELETE FROM prompts WHERE id=?", (rows[0]["id"],))
|
|
return rows
|
|
|
|
with patch.object(self.db, "_library_parents", side_effect=remove_anchor):
|
|
first = self.db.list_library(limit=1)
|
|
second = self.db.list_library(after=first["next_key"])
|
|
self.assertEqual([row["text"] for row in second["items"]], ["a"])
|
|
|
|
def test_options_preserve_exact_category_values_that_filters_accept(self):
|
|
self.save("Legacy", category=" Category ")
|
|
self.db.save_prompt_pair("Cat", "Noise", category=" Variant ")
|
|
options = self.db.get_library_options()["categories"]
|
|
self.assertEqual(options, [" Category ", " Variant "])
|
|
self.assertTrue(
|
|
all(self.db.list_library(category=value)["total"] == 1 for value in options)
|
|
)
|
|
|
|
def test_name_sort_uses_full_long_text_and_keeps_deleted_anchor_key(self):
|
|
prefix = "shared prefix " * 30
|
|
self.save(prefix + "Zulu")
|
|
alpha = self.save(prefix + "Alpha")
|
|
self.save(prefix + "Middle")
|
|
first = self.db.list_library(sort="name", limit=1)
|
|
self.assertEqual(first["items"][0]["id"], alpha)
|
|
self.assertEqual(first["next_key"]["value"], (prefix + "Alpha").lower())
|
|
with self.db.model.get_connection() as connection:
|
|
connection.execute("DELETE FROM prompts WHERE id=?", (alpha,))
|
|
rest = self.db.list_library(sort="name", after=first["next_key"])
|
|
self.assertEqual(
|
|
[item["text"] for item in rest["items"]],
|
|
[prefix + "Middle", prefix + "Zulu"],
|
|
)
|
|
|
|
def test_batch_previews_return_two_newest_links_per_parent_in_one_select(self):
|
|
parents = [self.save(f"Parent {index}") for index in range(200)]
|
|
expected = {}
|
|
for parent in parents[:3]:
|
|
ids = [
|
|
self.db.link_image_to_prompt(
|
|
parent, str(Path(self.temp.name) / f"{parent}-{index}.png")
|
|
)
|
|
for index in range(3)
|
|
]
|
|
expected[parent] = ids[-1:0:-1]
|
|
connection = self.db.model.get_connection()
|
|
statements = []
|
|
connection.set_trace_callback(statements.append)
|
|
try:
|
|
previews = self.db.get_library_previews(parents)
|
|
finally:
|
|
connection.set_trace_callback(None)
|
|
selects = [
|
|
query
|
|
for query in statements
|
|
if query.lstrip().upper().startswith(("SELECT", "WITH"))
|
|
]
|
|
self.assertEqual(len(selects), 1)
|
|
self.assertEqual(
|
|
{
|
|
parent: [image["id"] for image in images]
|
|
for parent, images in previews.items()
|
|
if images
|
|
},
|
|
expected,
|
|
)
|
|
self.assertEqual(self.db.get_library_previews([]), {})
|
|
with self.assertRaises(ValueError):
|
|
self.db.get_library_previews(parents + [9999])
|