Vito Sansevero
245d0cdc6a
- Thumbnail jobs and the legacy scan_output_dir listing leave library-removed files out (one mark check, mark_hides). - Missing-file links also forgets removal marks that hide nothing any more (file gone from a folder that exists, or changed), only under a readable, non-empty source, re-checked at apply. - The sidebar Images, Favorites and Needs review counts and the collection counts are read again when an image_delete job ends. - utils/media_folders.py: one skip for thumbnails and trash folders (.Trash, .Trash-<uid>, .Trashes, $RECYCLE.BIN) used by every walker, the image monitor and the delete review.
252 lines
10 KiB
Python
252 lines
10 KiB
Python
"""Removal-mark gaps after #480, and trash folders inside output roots (#501).
|
|
|
|
- Thumbnail jobs and the legacy auto-tag listing skip a library-removed file.
|
|
- Missing-file links also forgets the marks of files that are gone (only under
|
|
a source that can be read now).
|
|
- Every walker skips a trash folder that sits inside an output root.
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import shutil
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
try:
|
|
from tests.disk_waits import DISK_BOUND_SECONDS
|
|
from tests.test_image_delete_routes import ImageDeleteTestCase
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from disk_waits import DISK_BOUND_SECONDS
|
|
from test_image_delete_routes import ImageDeleteTestCase
|
|
|
|
from py.api import images as images_module
|
|
from py.api.admin import _collect_media_files
|
|
from py.api.autotag_routes import AutotagRoutesMixin
|
|
from utils.image_monitor import ImageGenerationHandler
|
|
|
|
MAINTENANCE_PREVIEW = "/prompt_manager/maintenance/preview"
|
|
MAINTENANCE_APPLY = "/prompt_manager/maintenance/apply"
|
|
|
|
|
|
class _GapCase(ImageDeleteTestCase):
|
|
def marks(self):
|
|
conn = sqlite3.connect(self.api.db.model.db_path)
|
|
try:
|
|
return [
|
|
row[0] for row in conn.execute("SELECT locator FROM image_removals")
|
|
]
|
|
finally:
|
|
conn.close()
|
|
|
|
def revision(self):
|
|
conn = sqlite3.connect(self.api.db.model.db_path)
|
|
try:
|
|
return conn.execute(
|
|
"SELECT revision FROM image_removals_meta WHERE singleton_id=1"
|
|
).fetchone()[0]
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
class TestThumbnailJobsSkipHiddenFiles(_GapCase):
|
|
async def test_missing_and_rebuild_plans_leave_a_removed_file_out(self):
|
|
self.write("hidden.png", b"hidden bytes", mtime=1_000_000)
|
|
self.write("shown.png", b"shown bytes", mtime=2_000_000)
|
|
await self.delete(["hidden.png"], "library")
|
|
thumbs = self.output_dir / "thumbnails"
|
|
for mode in ("missing", "rebuild"):
|
|
with self.subTest(mode=mode):
|
|
plan = self.api._plan_thumbnails(self.output_dir, thumbs, mode=mode)
|
|
sources = [src.name for src, _dst, _video in plan["targets"]]
|
|
self.assertEqual(sources, ["shown.png"])
|
|
self.assertEqual(plan["missing"], 1)
|
|
|
|
async def test_a_changed_file_gets_its_preview_again(self):
|
|
self.write("hidden.png", b"hidden bytes", mtime=1_000_000)
|
|
await self.delete(["hidden.png"], "library")
|
|
self.write("hidden.png", b"new and longer bytes", mtime=3_000_000)
|
|
plan = self.api._plan_thumbnails(
|
|
self.output_dir, self.output_dir / "thumbnails"
|
|
)
|
|
self.assertEqual([t[0].name for t in plan["targets"]], ["hidden.png"])
|
|
|
|
|
|
class TestLegacyAutotagListing(_GapCase):
|
|
async def test_scan_output_dir_leaves_a_removed_file_out(self):
|
|
self.write("hidden.png", b"hidden bytes", mtime=1_000_000)
|
|
self.write("shown.png", b"shown bytes", mtime=2_000_000)
|
|
await self.delete(["hidden.png"], "library")
|
|
with patch.object(
|
|
self.api, "_find_comfyui_output_dir", return_value=str(self.output_dir)
|
|
):
|
|
response = await self.client.get("/prompt_manager/scan_output_dir")
|
|
body = await response.json()
|
|
self.assertEqual([image["filename"] for image in body["images"]], ["shown.png"])
|
|
|
|
|
|
class TestMissingFileLinksPrunesMarks(_GapCase):
|
|
async def links_review(self, status=200):
|
|
return await self.post(
|
|
MAINTENANCE_PREVIEW, {"operation": "missing_links", "options": {}}, status
|
|
)
|
|
|
|
async def apply(self, token, status=202):
|
|
body = await self.post(MAINTENANCE_APPLY, {"token": token}, status)
|
|
if status != 202:
|
|
return body
|
|
job = self.api._jobs.get(body["job"]["id"])
|
|
await asyncio.wait_for(asyncio.shield(job.task), DISK_BOUND_SECONDS)
|
|
return job
|
|
|
|
async def removed(self, name):
|
|
self.write("other.png", b"another output") # the source lists something
|
|
self.write(name, f"bytes of {name}".encode(), mtime=1_000_000)
|
|
await self.delete([name], "library")
|
|
self.assertEqual(len(self.marks()), 1)
|
|
|
|
async def test_the_mark_of_a_gone_file_is_counted_and_forgotten(self):
|
|
await self.removed("gone.png")
|
|
# Not an unlink by PromptManager: the user deleted it elsewhere.
|
|
os.remove(self.output_dir / "gone.png")
|
|
plan = await self.links_review()
|
|
self.assertEqual(plan["scope"]["count"], 1)
|
|
self.assertEqual(
|
|
plan["scope"]["label"], "1 removed image whose file is gone or changed"
|
|
)
|
|
before = self.revision()
|
|
job = await self.apply(plan["token"])
|
|
self.assertEqual(job.state, "finished", job.summary())
|
|
self.assertEqual(self.marks(), [])
|
|
self.assertGreater(self.revision(), before)
|
|
|
|
async def test_the_mark_of_a_present_file_stays(self):
|
|
await self.removed("kept.png")
|
|
plan = await self.links_review()
|
|
self.assertEqual(plan["scope"]["count"], 0)
|
|
self.assertEqual(plan["scope"]["label"], "No links to missing files")
|
|
self.assertEqual(len(self.marks()), 1)
|
|
|
|
async def test_a_mark_under_an_unreadable_source_stays(self):
|
|
await self.removed("away.png")
|
|
os.remove(self.output_dir / "away.png")
|
|
with patch(
|
|
"database.maintenance.readable_roots",
|
|
side_effect=lambda r: [False] * len(r),
|
|
):
|
|
plan = await self.links_review()
|
|
self.assertEqual(plan["scope"]["count"], 0)
|
|
self.assertEqual(len(self.marks()), 1)
|
|
|
|
async def test_the_mark_of_a_changed_file_is_forgotten(self):
|
|
await self.removed("changed.png")
|
|
self.write("changed.png", b"other and longer bytes", mtime=4_000_000)
|
|
plan = await self.links_review()
|
|
self.assertEqual(plan["scope"]["count"], 1)
|
|
job = await self.apply(plan["token"])
|
|
self.assertEqual(job.state, "finished", job.summary())
|
|
self.assertEqual(self.marks(), [])
|
|
|
|
async def test_marks_under_an_empty_source_stay(self):
|
|
# An unplugged drive leaves its mount point: an empty, readable folder.
|
|
await self.removed("mounted.png")
|
|
for child in list(self.output_dir.iterdir()):
|
|
if child.is_dir():
|
|
shutil.rmtree(child)
|
|
else:
|
|
os.remove(child)
|
|
plan = await self.links_review()
|
|
self.assertEqual(plan["scope"]["count"], 0)
|
|
self.assertEqual(len(self.marks()), 1)
|
|
|
|
async def test_a_name_the_folder_still_lists_is_not_gone(self):
|
|
# A case-sensitive NTFS folder: lstat of the normcased locator fails.
|
|
await self.removed("Listed.png")
|
|
real_lstat = os.lstat
|
|
locator = self.marks()[0]
|
|
|
|
def lstat(path, *args, **kwargs):
|
|
if os.fspath(path) == locator:
|
|
raise FileNotFoundError(path)
|
|
return real_lstat(path, *args, **kwargs)
|
|
|
|
with (
|
|
patch("database.maintenance.os.lstat", side_effect=lstat),
|
|
patch("database.maintenance.os.path.normcase", side_effect=str.lower),
|
|
):
|
|
from database.maintenance import _stale_marks
|
|
|
|
stale = _stale_marks([(locator, 1, 1)], [str(self.output_dir)])
|
|
self.assertEqual(stale, [])
|
|
|
|
async def test_a_file_back_before_the_apply_changes_the_plan(self):
|
|
await self.removed("back.png")
|
|
path = self.output_dir / "back.png"
|
|
saved = Path(self.tmpdir) / "saved.png"
|
|
shutil.copy2(path, saved)
|
|
os.remove(path)
|
|
plan = await self.links_review()
|
|
shutil.copy2(saved, path)
|
|
job = await self.apply(plan["token"])
|
|
self.assertEqual(job.state, "failed", job.summary())
|
|
self.assertEqual(len(self.marks()), 1)
|
|
|
|
|
|
class TestWalkersSkipTrashFolders(_GapCase):
|
|
def seed(self):
|
|
self.write("kept.png", b"kept")
|
|
self.write(
|
|
"x.png", b"trashed x", folder=self.output_dir / ".Trash-1000" / "files"
|
|
)
|
|
self.write(
|
|
"y.png", b"trashed y", folder=self.output_dir / "sub" / "$RECYCLE.BIN"
|
|
)
|
|
self.write("z.png", b"trashed z", folder=self.output_dir / ".Trashes")
|
|
|
|
def test_walk_media_and_the_settled_entry(self):
|
|
self.seed()
|
|
walked = sorted(
|
|
Path(entry.path).name
|
|
for entry in images_module._walk_media(self.output_dir, 100)
|
|
)
|
|
self.assertEqual(walked, ["kept.png"])
|
|
for relative in (
|
|
os.path.join(".Trash-1000", "files", "x.png"),
|
|
os.path.join("sub", "$RECYCLE.BIN", "y.png"),
|
|
):
|
|
with self.subTest(relative=relative):
|
|
self.assertIsNone(
|
|
images_module._settled_entry(self.output_dir, relative)
|
|
)
|
|
|
|
def test_scan_duplicate_and_autotag_walkers(self):
|
|
self.seed()
|
|
collected = [
|
|
Path(p).name for p in _collect_media_files([self.output_dir], [".png"])
|
|
]
|
|
self.assertEqual(collected, ["kept.png"])
|
|
found, _truncated = AutotagRoutesMixin._collect_scan_files(self.output_dir)
|
|
self.assertEqual([p.name for p in found], ["kept.png"])
|
|
|
|
def test_the_image_monitor(self):
|
|
monitor = ImageGenerationHandler.__new__(ImageGenerationHandler)
|
|
monitor.supported_extensions = (".png",)
|
|
base = str(self.output_dir)
|
|
self.assertTrue(monitor.is_image_file(os.path.join(base, "a.png")))
|
|
for folder in (".Trash-1000", "$Recycle.Bin", ".Trashes", ".trash"):
|
|
with self.subTest(folder=folder):
|
|
self.assertFalse(
|
|
monitor.is_image_file(os.path.join(base, folder, "files", "a.png"))
|
|
)
|
|
|
|
async def test_the_workspace_listing(self):
|
|
self.seed()
|
|
self.assertEqual(await self.listed(), ["kept.png"])
|
|
|
|
def test_a_root_named_like_a_trash_folder_still_lists(self):
|
|
root = Path(self.tmpdir) / ".Trash-1000"
|
|
(root / "inner").mkdir(parents=True)
|
|
(root / "inner" / "a.png").write_bytes(b"a")
|
|
walked = [Path(e.path).name for e in images_module._walk_media(root, 10)]
|
|
self.assertEqual(walked, ["a.png"])
|