Files
ComfyAssets-ComfyUI_PromptM…/tests/test_image_removal.py
T
Vito 66ca988d2a feat(images): delete selected images, from the library or to the trash (#480) (#488)
* feat(trash): a freedesktop trash backend that renames without replacing or refuses (#480)

* feat(images): remove a file's records and hide it with a removal mark (#480)

Deletes every record of one file under every recorded spelling (image links,
tags, auto-tag marks, pending review items, category, collection memberships,
favorite) in one transaction. A Needs review item that is kept keeps its
annotation rows (ON DELETE RESTRICT). The image_removals table hides an exact
file (locator, size, mtime_ns) until it changes; linking skips a marked file,
and an explicit tag, category, favorite or collection write shows it again.

* feat(images): review and apply Delete images as a reviewed plan job (#480)

POST /prompt_manager/images/delete/preview reviews both modes; the
maintenance apply takes {token, mode}. Library mode removes the records and
writes the removal mark; trash mode moves the file first and removes its
records only after a successful move. The workspace listing, the output scan,
the duplicate find and Tag output images skip marked files.

* feat(images): Delete… in the Images selection bar, inspector and viewer (#480)

The concept's Remove links… place, D159: one review per selection that asks
each time (nothing preselected) whether to remove the files from the library
only or also move them to the trash. The browser test server uses a fake
trash.

* test(images): portable delete tests, the native no-trash proof, D159 in TESTING.md (#480)

* perf(images): cache the removal marks so a listing page reads no extra table (#480)

The marks reload only when this process writes or clears one, or a restore
replaces the database. The route snapshot gains the preview route.

* test(images): the full inspector specs expect D159's Delete… in the panel and the viewer (#480)

* fix(images): the delete review starts nothing without a chosen mode, and dims its primary until then (#480)

An aria-disabled primary ignores every activation, a chosen one starts once,
and the apply helper refuses a missing mode. Tests cover a forced click, a
double click, Enter and Space on the primary, Enter in the dialog and a
submit event, and the server refusing an apply without an exact mode.
2026-10-29 19:40:26 +00:00

320 lines
12 KiB
Python

"""Removing one file's records from the library, and removal marks (#480)."""
import hashlib
import os
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.image_removal_schema import (
image_removals_schema_error,
migrate_image_removals,
)
from database.models import PromptModel
from database.operations import PromptDatabase
def _locator(path):
return os.path.normcase(os.path.realpath(str(path)))
class _RemovalCase(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)
base = Path(self.temp.name)
self.root = base / "output"
self.root.mkdir()
# A second spelling of the folder (a symlink); Windows without the
# privilege has none, and the records then carry one spelling.
self.alias = base / "alias"
try:
self.alias.symlink_to(self.root, target_is_directory=True)
except OSError:
self.alias = None
self.db = PromptDatabase(str(base / "prompts.db"))
self.addCleanup(self.db.close_all)
self.epoch = self.db.get_execution_epoch()
def file(self, name, data=b"image-bytes"):
path = self.root / name
path.write_bytes(data)
return path
def removal(self, path):
root_key = os.path.normcase(os.path.realpath(str(self.root)))
return {
"spellings": [str(path)]
+ ([str(self.alias / path.name)] if self.alias else []),
"locator": _locator(path),
"favorites": [(root_key, path.name.lower())],
"category": (root_key, path.name.lower()),
}
def link(self, path, text="a prompt", spelled=None):
prompt_id = self.db.save_prompt(text)
image_id = self.db.link_image_to_prompt(str(prompt_id), str(spelled or path))
self.assertTrue(image_id)
return prompt_id, image_id
def tag(self, path, tags):
data = path.read_bytes()
info = os.stat(path)
return self.db.replace_image_tags(
locator=_locator(path),
sha256=hashlib.sha256(data).hexdigest(),
size=info.st_size,
mtime_ns=info.st_mtime_ns,
expected_epoch=self.epoch,
expected_tag_revision=0,
tags=tags,
)
def queue(self, path):
data = path.read_bytes()
info = os.stat(path)
return self.db.record_review_item(
locator=_locator(path),
spelled_path=os.path.realpath(str(path)),
sha256=hashlib.sha256(data).hexdigest(),
size=info.st_size,
mtime_ns=info.st_mtime_ns,
epoch=self.epoch,
resolver_version=1,
evidence={"candidates": []},
candidate_count=0,
limited=True,
)
def count(self, sql, *params):
conn = sqlite3.connect(self.db.model.db_path)
try:
return conn.execute(sql, params).fetchone()[0]
finally:
conn.close()
def decorate(self, path):
"""Give *path* one record of every kind."""
removal = self.removal(path)
root_key, relative = removal["favorites"][0]
self.tag(path, ["portrait"])
self.db.set_image_favorite(root_key, relative)
self.db.set_image_categories([removal["category"]], "keepers")
self.db.add_to_collection(items=[(root_key, relative)], name="Set")
return removal
class ImageRemovalRecordsTests(_RemovalCase):
def test_counts_and_removes_every_record_under_every_spelling(self):
path = self.file("a.png")
other = self.file("b.png", b"other")
orphan, _ = self.link(path, "only this image")
shared, _ = self.link(path, "two images")
self.db.link_image_to_prompt(str(shared), str(other))
# A record under the symlinked folder spelling (or a third prompt).
self.link(path, "aliased", spelled=(self.alias or self.root) / "a.png")
removal = self.decorate(path)
self.assertTrue(self.queue(path))
counted = self.db.count_removal_records([removal])
self.assertEqual(
counted["records"],
{
"prompt_links": 3,
"pair_links": 0,
"image_tags": 1,
"auto_tagged": 0,
"categories": 1,
"collections": 1,
"favorites": 1,
"review_items": 1,
},
)
self.assertEqual(counted["orphaned_prompts"]["count"], 2)
removed = self.db.remove_image_records(removal)
self.assertEqual(removed["prompt_links"], 3)
self.assertEqual(removed["review_items"], 1)
self.assertEqual(self.count("SELECT COUNT(*) FROM generated_images"), 1)
for table in (
"media_image_tags",
"media_annotation_files",
"media_review_items",
"image_favorites",
"image_categories",
"image_collection_items",
):
with self.subTest(table=table):
self.assertEqual(self.count(f"SELECT COUNT(*) FROM {table}"), 0)
# Prompts and the collection itself stay.
self.assertEqual(self.count("SELECT COUNT(*) FROM prompts"), 3)
self.assertEqual(self.count("SELECT COUNT(*) FROM image_collections"), 1)
self.assertTrue(path.exists())
self.assertIsNotNone(self.db.get_prompt_by_id(orphan))
def test_a_kept_review_item_keeps_the_annotation_rows(self):
path = self.file("a.png")
removal = self.decorate(path)
self.assertTrue(self.queue(path))
conn = sqlite3.connect(self.db.model.db_path)
try:
# A revised item is not settled; it still cites the bytes, whose
# foreign key is ON DELETE RESTRICT.
conn.execute("UPDATE media_review_items SET revision=1")
conn.commit()
finally:
conn.close()
removed = self.db.remove_image_records(removal)
self.assertEqual(removed["review_items"], 0)
self.assertEqual(removed["image_tags"], 1)
self.assertEqual(self.count("SELECT COUNT(*) FROM media_image_tags"), 0)
self.assertEqual(self.count("SELECT COUNT(*) FROM media_review_items"), 1)
self.assertEqual(self.count("SELECT COUNT(*) FROM media_annotation_files"), 1)
def test_mark_is_written_with_the_records_and_bumps_the_revision(self):
path = self.file("a.png")
self.link(path)
before = self.db.removal_revision()
info = os.stat(path)
self.db.remove_image_records(
self.removal(path), mark=(info.st_size, info.st_mtime_ns)
)
self.assertEqual(
self.db.removal_marks(),
{_locator(path): (info.st_size, info.st_mtime_ns)},
)
self.assertEqual(self.db.removal_revision(), before + 1)
self.assertEqual(self.db.clear_removal_marks([_locator(path)]), 1)
self.assertEqual(self.db.removal_marks(), {})
self.assertEqual(self.db.removal_revision(), before + 2)
self.assertEqual(self.db.clear_removal_marks([_locator(path)]), 0)
def test_a_failure_rolls_back_everything(self):
path = self.file("a.png")
self.link(path)
removal = self.decorate(path)
broken = {**removal, "category": ("only-one-value",)}
with self.assertRaises(sqlite3.Error):
self.db.remove_image_records(broken, mark=(1, 1))
self.assertEqual(self.count("SELECT COUNT(*) FROM generated_images"), 1)
self.assertEqual(self.count("SELECT COUNT(*) FROM media_image_tags"), 1)
self.assertEqual(self.db.removal_marks(), {})
def test_parity_with_the_legacy_single_delete(self):
path = self.file("a.png")
_prompt, image_id = self.link(path)
twin = PromptDatabase(str(Path(self.temp.name) / "twin.db"))
try:
twin_prompt = twin.save_prompt("a prompt")
twin_image = twin.link_image_to_prompt(str(twin_prompt), str(path))
self.assertTrue(twin.delete_image(twin_image))
self.db.remove_image_records(self.removal(path))
for table in (
"generated_images",
"image_prompt_pairs",
"workflow_documents",
):
with self.subTest(table=table):
conn = sqlite3.connect(twin.model.db_path)
try:
legacy = conn.execute(
f"SELECT COUNT(*) FROM {table}"
).fetchone()[0]
finally:
conn.close()
self.assertEqual(
self.count(f"SELECT COUNT(*) FROM {table}"), legacy
)
finally:
twin.close_all()
self.assertIsNotNone(image_id)
class RemovalMarkGateTests(_RemovalCase):
"""A marked file is never linked again; any record write shows it again."""
def hide(self, path):
info = os.stat(path)
self.db.remove_image_records(
self.removal(path), mark=(info.st_size, info.st_mtime_ns)
)
def test_a_marked_file_is_not_linked_again(self):
path = self.file("a.png")
self.hide(path)
prompt_id = self.db.save_prompt("scanned again")
self.assertEqual(self.db.link_image_to_prompt(str(prompt_id), str(path)), 0)
self.assertEqual(self.count("SELECT COUNT(*) FROM generated_images"), 0)
self.assertIn(_locator(path), self.db.removal_marks())
def test_a_changed_file_drops_its_stale_mark_and_links(self):
path = self.file("a.png")
self.hide(path)
path.write_bytes(b"new bytes, new size")
prompt_id = self.db.save_prompt("a new image at the same path")
self.assertTrue(self.db.link_image_to_prompt(str(prompt_id), str(path)))
self.assertEqual(self.db.removal_marks(), {})
def test_explicit_record_writes_clear_the_mark(self):
writers = {
"favorite": lambda path, key: self.db.set_image_favorite(*key),
"category": lambda path, key: self.db.set_image_categories([key], "c"),
"collection": lambda path, key: self.db.add_to_collection(
items=[key], name="Set"
),
"tags": lambda path, key: self.tag(path, ["t"]),
}
for kind, write in writers.items():
with self.subTest(kind=kind):
path = self.file(f"{kind}.png", kind.encode())
self.hide(path)
self.assertIn(_locator(path), self.db.removal_marks())
write(path, self.removal(path)["favorites"][0])
self.assertNotIn(_locator(path), self.db.removal_marks())
class ImageRemovalSchemaTests(unittest.TestCase):
def setUp(self):
self.conn = sqlite3.connect(":memory:")
self.addCleanup(self.conn.close)
def test_missing_is_allowed_only_for_old_backups(self):
self.assertIsNone(image_removals_schema_error(self.conn))
self.assertIsNotNone(
image_removals_schema_error(self.conn, allow_missing=False)
)
def test_install_is_valid_and_idempotent(self):
migrate_image_removals(self.conn)
migrate_image_removals(self.conn)
self.assertIsNone(image_removals_schema_error(self.conn, allow_missing=False))
self.assertEqual(
self.conn.execute("SELECT revision FROM image_removals_meta").fetchall(),
[(0,)],
)
def test_invalid_values_and_objects_are_refused(self):
migrate_image_removals(self.conn)
self.conn.execute("INSERT INTO image_removals VALUES ('', 1, 1, 'now')")
self.assertIn("invalid values", image_removals_schema_error(self.conn))
self.conn.execute("DELETE FROM image_removals")
self.conn.execute("CREATE INDEX extra ON image_removals(size)")
self.assertIn("unsupported", image_removals_schema_error(self.conn))
def test_a_foreign_table_definition_is_refused(self):
self.conn.execute("CREATE TABLE image_removals (locator TEXT)")
with self.assertRaises(sqlite3.DatabaseError):
migrate_image_removals(self.conn)
if __name__ == "__main__":
unittest.main()