224 lines
9.3 KiB
Python
224 lines
9.3 KiB
Python
"""Collections (#331): the database operations."""
|
|
|
|
import sqlite3
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
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
|
|
|
|
try: # guarded so a missing module fails assertions, not test loading
|
|
from database.image_collections import (
|
|
CollectionFilesMissing,
|
|
CollectionInvalid,
|
|
CollectionMissing,
|
|
CollectionNameTaken,
|
|
READ_SQL,
|
|
)
|
|
except ImportError:
|
|
CollectionFilesMissing = CollectionMissing = CollectionNameTaken = None
|
|
CollectionInvalid = ValueError
|
|
READ_SQL = None
|
|
|
|
ROOT = "/real/output"
|
|
OTHER = "/real/other"
|
|
|
|
|
|
class ImageCollectionOperationsTests(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.db = PromptDatabase(str(Path(self.temp.name) / "prompts.db"))
|
|
self.addCleanup(self.db.close_all)
|
|
|
|
def require(self):
|
|
self.assertTrue(
|
|
hasattr(self.db, "add_to_collection"),
|
|
"ImageCollectionOperations is not implemented",
|
|
)
|
|
|
|
def add(self, items, **target):
|
|
return self.db.add_to_collection(items=items, **target)
|
|
|
|
def test_add_creates_a_collection_by_name(self):
|
|
self.require()
|
|
result = self.add(
|
|
[(ROOT, "a.png"), (ROOT, "b.png"), (OTHER, "a.png")], name="LoRa testing"
|
|
)
|
|
collection = result["collection"]
|
|
self.assertEqual(collection["name"], "LoRa testing")
|
|
self.assertEqual(
|
|
result,
|
|
{"collection": collection, "created": True, "added": 3, "already": 0},
|
|
)
|
|
self.assertEqual(self.db.list_collections(), [collection])
|
|
self.assertTrue(self.db.collection_exists(collection["id"]))
|
|
self.assertEqual(self.db.get_collection(collection["id"]), collection)
|
|
self.assertEqual(
|
|
self.db.collection_members(collection["id"], ROOT), {"a.png", "b.png"}
|
|
)
|
|
self.assertEqual(
|
|
self.db.root_collection_members(OTHER), {collection["id"]: {"a.png"}}
|
|
)
|
|
self.assertEqual(self.db.collection_counts(), {collection["id"]: 3})
|
|
self.assertEqual(self.db.file_collections(ROOT, "a.png"), [collection["id"]])
|
|
self.assertEqual(self.db.file_collections(ROOT, "zzz.png"), [])
|
|
|
|
def test_add_by_name_ignores_case_and_unicode_form(self):
|
|
self.require()
|
|
first = self.add([(ROOT, "a.png")], name="Café grids")["collection"]
|
|
# NFD spelling (macOS) in another case joins the same collection.
|
|
second = self.add([(ROOT, "b.png")], name="CAFÉ GRIDS")
|
|
self.assertEqual(second["collection"], first)
|
|
self.assertFalse(second["created"])
|
|
self.assertEqual(len(self.db.list_collections()), 1)
|
|
self.assertEqual(
|
|
self.db.collection_members(first["id"], ROOT), {"a.png", "b.png"}
|
|
)
|
|
|
|
def test_add_counts_already_present(self):
|
|
self.require()
|
|
made = self.add([(ROOT, "a.png"), (ROOT, "b.png")], name="Set")
|
|
result = self.add(
|
|
[(ROOT, "b.png"), (ROOT, "c.png"), (ROOT, "c.png")],
|
|
collection_id=made["collection"]["id"],
|
|
)
|
|
self.assertEqual(
|
|
(result["created"], result["added"], result["already"]), (False, 1, 1)
|
|
)
|
|
self.assertEqual(self.db.collection_counts(), {made["collection"]["id"]: 3})
|
|
|
|
def test_add_with_a_new_name_and_nothing_inserted_creates_nothing(self):
|
|
self.require()
|
|
with self.assertRaises(CollectionFilesMissing):
|
|
self.add([], name="Empty")
|
|
self.assertEqual(self.db.list_collections(), [])
|
|
# An existing name with nothing to add is not an error.
|
|
made = self.add([(ROOT, "a.png")], name="Kept")["collection"]
|
|
self.assertEqual(
|
|
self.add([], name="kept"),
|
|
{"collection": made, "created": False, "added": 0, "already": 0},
|
|
)
|
|
|
|
def test_add_to_a_deleted_collection_raises_missing(self):
|
|
self.require()
|
|
made = self.add([(ROOT, "a.png")], name="Gone")["collection"]
|
|
self.db.delete_collection(made["id"])
|
|
with self.assertRaises(CollectionMissing):
|
|
self.add([(ROOT, "b.png")], collection_id=made["id"])
|
|
with self.assertRaises(CollectionMissing):
|
|
self.db.remove_from_collection(made["id"], [(ROOT, "a.png")])
|
|
with self.assertRaises(CollectionMissing):
|
|
self.db.rename_collection(made["id"], "Back")
|
|
self.assertFalse(self.db.collection_exists(made["id"]))
|
|
self.assertIsNone(self.db.get_collection(made["id"]))
|
|
|
|
def test_remove_returns_the_count(self):
|
|
self.require()
|
|
made = self.add([(ROOT, "a.png"), (ROOT, "b.png")], name="Set")["collection"]
|
|
result = self.db.remove_from_collection(
|
|
made["id"], [(ROOT, "a.png"), (ROOT, "a.png"), (ROOT, "never.png")]
|
|
)
|
|
self.assertEqual(result, {"collection": made, "removed": 1})
|
|
self.assertEqual(self.db.collection_members(made["id"], ROOT), {"b.png"})
|
|
# Removing the last member keeps the (now empty) collection.
|
|
self.db.remove_from_collection(made["id"], [(ROOT, "b.png")])
|
|
self.assertEqual(self.db.list_collections(), [made])
|
|
self.assertEqual(self.db.collection_counts(), {made["id"]: 0})
|
|
|
|
def test_rename_refuses_a_taken_name(self):
|
|
self.require()
|
|
first = self.add([(ROOT, "a.png")], name="First")["collection"]
|
|
second = self.add([(ROOT, "a.png")], name="Second")["collection"]
|
|
with self.assertRaises(CollectionNameTaken) as caught:
|
|
self.db.rename_collection(second["id"], " first ")
|
|
self.assertEqual(caught.exception.collection_id, first["id"])
|
|
self.assertEqual(self.db.get_collection(second["id"]), second)
|
|
|
|
def test_rename_accepts_a_case_variant_of_its_own_name(self):
|
|
self.require()
|
|
made = self.add([(ROOT, "a.png")], name="lora grids")["collection"]
|
|
renamed = self.db.rename_collection(made["id"], "LoRA grids ")
|
|
self.assertEqual(renamed, {"id": made["id"], "name": "LoRA grids"})
|
|
self.assertEqual(self.db.list_collections(), [renamed])
|
|
other = self.db.rename_collection(made["id"], "Portraits")
|
|
self.assertEqual(other["name"], "Portraits")
|
|
|
|
def test_delete_unknown_raises_missing(self):
|
|
self.require()
|
|
made = self.add([(ROOT, "a.png")], name="Doomed")["collection"]
|
|
self.assertTrue(self.db.delete_collection(made["id"]))
|
|
self.assertEqual(self.db.root_collection_members(ROOT), {})
|
|
with self.assertRaises(CollectionMissing):
|
|
self.db.delete_collection(made["id"])
|
|
with self.assertRaises(CollectionMissing):
|
|
self.db.delete_collection(999)
|
|
|
|
def test_list_is_in_name_key_order(self):
|
|
self.require()
|
|
for name in ("beta", "Alpha", "été", "Zeta", "alpha two"):
|
|
self.add([(ROOT, "a.png")], name=name)
|
|
self.assertEqual(
|
|
[row["name"] for row in self.db.list_collections()],
|
|
["Alpha", "alpha two", "beta", "Zeta", "été"],
|
|
)
|
|
|
|
def test_names_are_trimmed_and_collapsed(self):
|
|
self.require()
|
|
made = self.add([(ROOT, "a.png")], name=" LoRa \t testing\n ")["collection"]
|
|
self.assertEqual(made["name"], "LoRa testing")
|
|
for bad in ("", " ", "x" * 101, "bell\x07", 5, None):
|
|
with self.subTest(name=bad):
|
|
with self.assertRaises(CollectionInvalid):
|
|
self.add([(ROOT, "a.png")], name=bad)
|
|
self.assertEqual(len(self.db.list_collections()), 1)
|
|
|
|
def test_targets_and_items_are_validated(self):
|
|
self.require()
|
|
made = self.add([(ROOT, "a.png")], name="Set")["collection"]
|
|
for target in (
|
|
{},
|
|
{"name": "x", "collection_id": made["id"]},
|
|
{"collection_id": True},
|
|
{"collection_id": 0},
|
|
{"collection_id": "1"},
|
|
):
|
|
with self.subTest(target=target):
|
|
with self.assertRaises(CollectionInvalid):
|
|
self.add([(ROOT, "a.png")], **target)
|
|
for items in ([("root",)], [(ROOT, 5)], "a.png", [(ROOT, "x" * 4097)]):
|
|
with self.subTest(items=items):
|
|
with self.assertRaises(CollectionInvalid):
|
|
self.add(items, collection_id=made["id"])
|
|
|
|
def test_reads_use_an_index(self):
|
|
self.require()
|
|
self.add([(ROOT, "a.png")], name="Set")
|
|
conn = sqlite3.connect(self.db.model.db_path)
|
|
try:
|
|
for sql in READ_SQL:
|
|
plan = " ".join(
|
|
row[3]
|
|
for row in conn.execute(
|
|
"EXPLAIN QUERY PLAN " + sql, [None] * sql.count("?")
|
|
)
|
|
)
|
|
with self.subTest(sql=sql):
|
|
self.assertNotIn("SCAN image_collection_items", plan)
|
|
if "?" in sql: # a lookup, not the whole list
|
|
self.assertNotIn("SCAN image_collections", plan)
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|