nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
3f085d0724778f07dffd8d3eff8bea13d9c08353
ComfyAssets-ComfyUI_PromptM…/tests/test_api_duplicate_files.py
T
Vito d0339128e4 test: isolate API fixture databases from other suites; CI lint pins, failing bandit hook, npm cache (#346, #352, #409) (#411) 
* ci: pin black and flake8, let the bandit hook fail, cache npm in the JS job (#352)

The lint job installs black==26.10.1 and flake8==7.4.1 from
requirements-dev.txt (the versions CI resolved on main today), so a new
release cannot turn CI red with no code change.

The local bandit hook ran 'bash -c ... "$@"; true': it could never
fail, and bash bound the first staged file to $0 so it was never
scanned. It now runs 'bandit -lll -q' directly, failing on high
severity exactly like code-quality.yml.

The frontend modules job's setup-node caches npm keyed on
tools/vendor/package-lock.json.

* test(history): wait on the writer's state, not short wall-clock pauses (#352)

Four 0.08 s pauses before asserting that the worker had not probed yet
now wait until the idle worker has run a whole loop pass at the current
manual clock value (three clock reads), so a slow runner cannot end the
pause before the worker decides. A mutation that probes 20 ms early now
fails two tests deterministically.

The real-clock bounded close no longer asserts close(0.02) returned in
under 0.2 s wall time; it asserts the in-flight write is still blocked
and the worker alive, like the output writer test. The manual-clock
bound stays exact. The last 0.1 s non-blocking submit check uses the
1 s NON_BLOCKING_CEILING from #351.

* test(api): give each API fixture database its own temp folder (#346)

Every PromptManagerAPI sweeps its database's folder at startup for the
stages, uploads and safety-backup copies an interrupted restore left.
The API test fixtures created their database directly in the system
temp folder, so any other test process starting an API there (another
worktree's suite, a pre-push hook, an agent's targeted run) deleted
this process's in-flight restore files. The restore, epoch and
retirement tests then failed with 'staged restore changed', 'unable to
open database file' or 503 instead of 409. Busy machines fail because
they are running other suites, not because they are slow: the same
CPU load with a private TMPDIR passes.

private_temp_db() puts the fixture database in a new folder removed at
cleanup (with any safety backups a restore left beside it). A
regression test sweeps a stand-in shared folder in the middle of a
multipart restore, as another suite's startup would.

* test(history): count worker parks, hold saturation blocks until released (#352)

wait_for_worker_pass now counts the idle worker's parks on its wake
event (two parks from now = one whole pass begun after the call)
instead of assuming a fixed number of clock reads per pass, which a
future clock read would have made vacuous. The real-clock bounded
close keeps a loose upper bound (0.02 s + NON_BLOCKING_CEILING).

The two saturation tests (history and output writer) blocked the
worker with release.wait(3) while they queued 1023 more items. On a
loaded runner the fill took longer than 3 s, the worker unblocked on
its own and drained part of the queue (pending 868..1017 != 1024),
and cleanup then reported a leaked worker. The block now holds until
the test releases it; 60 s is only a hang guard, and cleanup always
sets the release. Under one core with 8 busy loops: 8/8 runs failed
before, 6/6 pass after.

* test(maintenance): wait on duplicate jobs with a disk-bound hang guard (#409)

On Windows py3.13 CI (PR #405) a one-file duplicate cleanup missed its
10 s WAIT. The test already waits on the job's own task; the 10 s bound
acted as a deadline on disk-bound work. In the same 50 s window the
following tests' setups took up to 9 s against 0.1-0.6 s around it
while the test bodies stayed fast, which is consistent with a
process-wide disk I/O stall on the runner rather than a lease, handle,
executor or job stall (the log has no thread dump to show more).

tests/disk_waits.py holds the one disk-bound hang guard (90 s per wait,
under the 120 s per-test watchdog), shared with RESTORE_SECONDS (#382).
Waits on find and cleanup job tasks in the duplicate-files tests use
it; gates and server-state polls keep WAIT. The keeper-validation test
turns the safety backup off, since it does not test backups.
2026-10-28 05:18:00 +00:00

1002 lines
43 KiB
Python
Raw Blame History

"""Duplicate image files on the job registry (#249 System S6, Task 17).
``POST /prompt_manager/jobs/duplicates`` finds content-identical media files
as a ``read`` Activity job; ``GET /prompt_manager/maintenance/duplicates/{id}``
pages the groups it holds; ``duplicate_files`` previews a removal bound to the
reviewed keepers, and its apply job removes only the reviewed, revalidated
non-keepers, moving their database links to the kept copy. The legacy gallery
routes keep their shapes on the same registry kinds.
"""
import asyncio
import hashlib
import json
import os
import sqlite3
import threading
from pathlib import Path
from unittest.mock import patch
try:
from tests.disk_waits import DISK_BOUND_SECONDS
from tests.open_handles import open_handles_under
from tests.test_api_maintenance import ApplyTestCase, _digest
except ImportError: # discovered with tests/ as the top-level directory
from disk_waits import DISK_BOUND_SECONDS
from open_handles import open_handles_under
from test_api_maintenance import ApplyTestCase, _digest
from py.api import duplicate_files # noqa: E402
from py.config import GalleryConfig, PromptManagerConfig # noqa: E402
# Gates and server-state polls. Waits on a find or cleanup job's task use
# DISK_BOUND_SECONDS: those jobs hash, back up and unlink files (#409).
WAIT = 10
FIND = "/prompt_manager/jobs/duplicates"
GROUPS = "/prompt_manager/maintenance/duplicates/"
PREVIEW = "/prompt_manager/maintenance/preview"
STREAM = "/prompt_manager/scan_duplicates/stream"
STATUS = "/prompt_manager/scan_duplicates/status"
LEGACY_DELETE = "/prompt_manager/delete_duplicate_images"
def _sse_payloads(body):
return [
json.loads(line[len("data: ") :])
for line in body.splitlines()
if line.startswith("data: ")
]
def _sha(path):
with open(path, "rb") as handle:
return hashlib.sha256(handle.read()).hexdigest()
class DuplicateFilesTestCase(ApplyTestCase):
async def get_application(self):
app = await super().get_application()
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
return app
def write(self, name, data, *, folder=None, mtime=None):
path = Path(folder or self.output_dir) / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(data)
if mtime is not None:
os.utime(path, (mtime, mtime))
return path
def second_source(self):
"""A second valid source outside ComfyUI (an extra gallery root)."""
second = self.outside_dir / "renders"
second.mkdir(parents=True, exist_ok=True)
os.environ["PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"] = str(second)
return second
def trio(self):
"""a.png (oldest), b.png and c.png share bytes; d.png is unique."""
self.write("a.png", b"same bytes", mtime=1_000_000)
self.write("b.png", b"same bytes", mtime=2_000_000)
self.write("sub/c.png", b"same bytes", mtime=3_000_000)
self.write("d.png", b"unique", mtime=4_000_000)
async def find(self):
response = await self.client.request("POST", FIND)
text = await response.text()
self.assertEqual(response.status, 202, text)
self.assertNotIn(self.tmpdir, text)
job = self.api._jobs.get(json.loads(text)["job"]["id"])
await asyncio.wait_for(asyncio.shield(job.task), DISK_BOUND_SECONDS)
return job
async def groups(self, job_id, status=200, **query):
params = "&".join(f"{key}={value}" for key, value in query.items())
response = await self.client.request("GET", f"{GROUPS}{job_id}?{params}")
text = await response.text()
self.assertEqual(response.status, status, text)
self.assertNotIn(self.tmpdir, text)
return json.loads(text)
async def review(self, job_id, keep, status=200):
return await self.post(
PREVIEW,
{
"operation": "duplicate_files",
"options": {"job_id": job_id, "keep": keep},
},
status,
)
async def found_group(self):
job = await self.find()
[group] = (await self.groups(job.id))["groups"]
return job, group
@staticmethod
def keep_of(group, path, root_index=None):
[entry] = [
f
for f in group["files"]
if f["path"] == path and root_index in (None, f["root_index"])
]
return {
"group_id": group["group_id"],
"root_index": entry["root_index"],
"path": entry["path"],
}
async def remove_all_but(self, path):
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, path)])
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
return plan, cleanup
async def until(self, predicate):
"""Wait for a server-side state (a frame at a time), never a fixed timer."""
deadline = asyncio.get_running_loop().time() + WAIT
while not predicate():
if asyncio.get_running_loop().time() > deadline:
self.fail("the expected server state never appeared")
await asyncio.sleep(0)
def gate_hashing(self):
gate = threading.Event()
calls = []
original = self.api._calculate_file_hash
def gated(path):
calls.append(path)
gate.wait(WAIT)
return original(path)
self.api._calculate_file_hash = gated
self.addCleanup(gate.set)
return gate, calls
class TestFindJob(DuplicateFilesTestCase):
async def test_find_is_listed_in_activity_and_legacy_stream_attaches(self):
self.trio()
gate, calls = self.gate_hashing()
response = await self.client.request("POST", FIND)
self.assertEqual(response.status, 202)
summary = (await response.json())["job"]
listed = (
await (await self.client.request("GET", "/prompt_manager/jobs")).json()
)["jobs"]
self.assertEqual([j["id"] for j in listed], [summary["id"]])
self.assertEqual(summary["kind"], "duplicates")
self.assertEqual(summary["title"], "Find duplicate image files")
self.assertEqual(summary["admission"], "read")
self.assertTrue(summary["cancellable"])
self.assertEqual(summary["counters"], {"groups_found": 0})
# The legacy gallery attaches to the same job, and its status says so.
legacy = asyncio.ensure_future(self.client.request("POST", STREAM))
job = self.api._jobs.get(summary["id"])
await self.until(lambda: job.subscriber_count > 0)
status = await (await self.client.request("GET", STATUS)).json()
self.assertTrue(status["running"])
gate.set()
body = await (await legacy).text()
job = self.api._jobs.get(summary["id"])
await asyncio.wait_for(asyncio.shield(job.task), DISK_BOUND_SECONDS)
events = _sse_payloads(body)
done = events[-1]
self.assertEqual(done["type"], "complete")
self.assertEqual(done["found"], 1)
self.assertEqual(
sorted(i["filename"] for i in done["duplicates"][0]["images"]),
["a.png", "b.png", "c.png"],
)
self.assertEqual(len(calls), 4) # one pass: each file hashed once
self.assertEqual(job.state, "finished")
self.assertEqual(job.counters, {"groups_found": 1})
[row] = job.outcomes["groups_found"]
self.assertEqual(row["path"], "a.png")
self.assertEqual(row["root_index"], 0)
self.assertEqual(row["detail"], "3 identical files · 10 B each")
self.assertEqual(job.status_line, "Found 1 group of duplicate files")
finished = await (await self.client.request("GET", STATUS)).json()
self.assertFalse(finished["running"])
self.assertEqual(finished["last_event"]["found"], 1)
async def test_legacy_stream_starts_a_registry_job(self):
self.trio()
body = await (await self.client.request("POST", STREAM)).text()
self.assertEqual(_sse_payloads(body)[-1]["type"], "complete")
[job] = self.api._jobs.list()
self.assertEqual(job.kind, "duplicates")
self.assertEqual(job.title, "Find duplicate image files (gallery page)")
self.assertNotIn(self.tmpdir, body)
# Its groups are reviewable in the workspace too.
page = await self.groups(job.id)
self.assertEqual(page["total"], 1)
async def test_find_covers_every_source_with_its_index(self):
second = self.second_source()
self.write("x.png", b"twin", folder=second)
GalleryConfig.MONITORING_DIRECTORIES = [
self._root(self.output_dir),
self._root(second),
]
self.write("x.png", b"twin")
job = await self.find()
[group] = (await self.groups(job.id))["groups"]
self.assertEqual(
sorted((f["root_index"], f["path"]) for f in group["files"]),
[(0, "x.png"), (1, "x.png")],
)
self.assertEqual({f["root"] for f in group["files"]}, {"output", "renders"})
async def test_find_can_be_cancelled(self):
self.trio()
gate, _calls = self.gate_hashing()
response = await self.client.request("POST", FIND)
job_id = (await response.json())["job"]["id"]
stream = asyncio.ensure_future(self.client.request("POST", STREAM))
found = self.api._jobs.get(job_id)
await self.until(lambda: found.subscriber_count > 0)
cancel = await self.client.request(
"POST", f"/prompt_manager/jobs/{job_id}/cancel"
)
self.assertTrue((await cancel.json())["cancellation_requested"])
gate.set()
job = self.api._jobs.get(job_id)
await asyncio.wait_for(asyncio.shield(job.task), DISK_BOUND_SECONDS)
self.assertEqual(job.state, "cancelled")
last = _sse_payloads(await (await stream).text())[-1]
self.assertEqual(last, {"type": "error", "message": "Duplicate scan cancelled"})
await self.groups(job_id, 409)
async def test_second_find_is_busy_and_names_the_running_job(self):
self.trio()
gate, _calls = self.gate_hashing()
first = (await (await self.client.request("POST", FIND)).json())["job"]
second = await self.client.request("POST", FIND)
self.assertEqual(second.status, 409)
body = await second.json()
self.assertEqual(body["code"], "busy")
self.assertEqual(body["running"]["id"], first["id"])
gate.set()
async def test_failed_find_is_coded_and_releases_the_kind(self):
self.write("a.png", b"x")
with patch.object(
self.api, "_hash_media_groups", side_effect=OSError(f"boom {self.tmpdir}")
):
with self.assertLogs("prompt_manager", "ERROR"):
job = await self.find()
self.assertEqual(job.state, "failed")
self.assertEqual(job.status_line, "The duplicate search failed.")
self.assertNotIn("boom", json.dumps(job.report()))
await self.groups(job.id, 409)
self.assertEqual((await self.find()).state, "finished")
class TestGroupPages(DuplicateFilesTestCase):
async def test_groups_are_paged_ten_at_a_time(self):
for index in range(12):
for copy in range(2):
self.write(f"g{index:02d}_{copy}.png", f"group {index}".encode())
job = await self.find()
first = await self.groups(job.id)
self.assertEqual((first["total"], first["offset"], first["limit"]), (12, 0, 10))
self.assertEqual(len(first["groups"]), 10)
rest = await self.groups(job.id, offset=10, limit=10)
self.assertEqual(len(rest["groups"]), 2)
ids = [g["group_id"] for g in first["groups"] + rest["groups"]]
self.assertEqual(len(set(ids)), 12)
entry = first["groups"][0]["files"][0]
for key in ("path", "root", "root_index", "url", "size", "media_type"):
self.assertIn(key, entry)
self.assertTrue(entry["url"].endswith("?root=0"))
for query in ({"limit": 11}, {"limit": 0}, {"offset": -1}, {"offset": "x"}):
with self.subTest(query=query):
reply = await self.groups(job.id, 400, **query)
self.assertEqual(reply["code"], "invalid-request")
async def test_unknown_or_running_job_has_no_groups(self):
reply = await self.groups("f" * 32, 404)
self.assertEqual(reply["code"], "not-found")
self.trio()
gate, _calls = self.gate_hashing()
job_id = (await (await self.client.request("POST", FIND)).json())["job"]["id"]
reply = await self.groups(job_id, 409)
self.assertEqual(reply["code"], "job-running")
gate.set()
class TestReview(DuplicateFilesTestCase):
async def test_preview_lists_only_non_keepers_and_writes_nothing(self):
self.trio()
prompt = self.save("Linked")
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "b.png"))
job, group = await self.found_group()
db_path = self.api.db.model.db_path
files = [db_path, db_path + "-wal"]
before = [_digest(path) for path in files]
media = {p: _digest(p) for p in self.output_dir.rglob("*") if p.is_file()}
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.assertEqual([_digest(path) for path in files], before)
self.assertEqual(
{p: _digest(p) for p in self.output_dir.rglob("*") if p.is_file()},
media,
)
self.assertEqual(plan["operation"], "duplicate_files")
self.assertEqual(
plan["scope"], {"count": 2, "label": "2 duplicate files in 1 group"}
)
self.assertTrue(plan["media"]["deletes_files"])
self.assertEqual(
plan["media"]["label"],
"Deletes 2 files (20 B) from disk; the kept copies stay. "
"The safety backup covers the database, not these files.",
)
self.assertEqual(
[row["label"] for row in plan["details"]], ["b.png", "sub/c.png"]
)
self.assertEqual(
plan["details"][0]["detail"],
"output · a copy of a.png · 1 database link moves to the kept copy",
)
self.assertEqual(plan["protected"], {"count": 0, "reasons": []})
async def test_cannot_delete_keeper_or_outside_group(self):
# About which files may go, not the backup: skip its fsync-heavy copy.
PromptManagerConfig.SAFETY_BACKUP = False
self.trio()
self.write("e.png", b"other pair")
self.write("f.png", b"other pair")
job = await self.find()
groups = (await self.groups(job.id))["groups"]
trio = next(g for g in groups if len(g["files"]) == 3)
pair = next(g for g in groups if len(g["files"]) == 2)
bad = [
# a file of another group named as this group's keeper
[{**self.keep_of(trio, "a.png"), "path": "e.png"}],
# a file that is in no group
[{**self.keep_of(trio, "a.png"), "path": "d.png"}],
# the right path under another source index
[{**self.keep_of(trio, "a.png"), "root_index": 1}],
[{**self.keep_of(trio, "a.png"), "group_id": "0" * 64}],
[{**self.keep_of(trio, "a.png"), "extra": 1}],
[],
"a.png",
]
for keep in bad:
with self.subTest(keep=keep):
reply = await self.review(job.id, keep, 400)
self.assertEqual(reply["code"], "invalid-request")
# Two keepers of one group: only the third copy goes.
plan = await self.review(
job.id,
[self.keep_of(trio, "a.png"), self.keep_of(trio, "sub/c.png")],
)
self.assertEqual([row["label"] for row in plan["details"]], ["b.png"])
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
self.assertEqual(cleanup.counters["removed"], 1)
self.assertTrue((self.output_dir / "a.png").is_file())
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
self.assertFalse((self.output_dir / "b.png").exists())
# The untouched group is untouched.
self.assertTrue((self.output_dir / "e.png").is_file())
self.assertTrue((self.output_dir / "f.png").is_file())
self.assertEqual(len(pair["files"]), 2)
async def test_a_source_removed_since_the_find_is_refused(self):
self.trio()
job, group = await self.found_group()
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.second_source())]
reply = await self.review(job.id, [self.keep_of(group, "a.png")], 409)
self.assertEqual(reply["code"], "sources-changed")
def test_containment_is_decided_as_text_before_the_filesystem(self):
root = str(self.output_dir)
outside = os.path.join(str(self.outside_dir), "x.png")
touched = []
real_lstat = os.lstat
def lstat(path, *args, **kwargs):
touched.append(str(path))
return real_lstat(path, *args, **kwargs)
with patch.object(duplicate_files.os, "lstat", side_effect=lstat):
for candidate in (
outside,
os.path.join(root, "..", "outside", "x.png"),
os.path.join(root, "thumbnails", "a_thumb.png"),
os.path.join(root, "notes.txt"),
):
with self.subTest(candidate=candidate):
reason = duplicate_files.file_problem(candidate, root, "0" * 64, 1)
self.assertEqual(reason, duplicate_files.OUTSIDE_SOURCE)
self.assertEqual(touched, [])
@staticmethod
def _can_symlink(directory):
probe = Path(directory) / "probe-link"
try:
os.symlink(str(Path(directory) / "x"), str(probe))
except (OSError, NotImplementedError):
return False
probe.unlink()
return True
async def test_symlink_rejected(self):
if not self._can_symlink(self.tmpdir): # pragma: no cover - Windows w/o rights
self.skipTest("symlinks are not available")
self.trio()
job, group = await self.found_group()
# b.png becomes a link to the keeper after the find: never removed.
(self.output_dir / "b.png").unlink()
os.symlink(str(self.output_dir / "a.png"), str(self.output_dir / "b.png"))
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.assertEqual([row["label"] for row in plan["details"]], ["sub/c.png"])
self.assertEqual(
plan["protected"],
{"count": 1, "reasons": ["1 file is a link, not a regular file"]},
)
# A keeper that became a link protects its whole group.
(self.output_dir / "a.png").rename(self.output_dir / "real.png")
os.symlink(str(self.output_dir / "real.png"), str(self.output_dir / "a.png"))
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.assertEqual(plan["scope"]["count"], 0)
self.assertEqual(
plan["protected"]["reasons"],
["2 files kept: their kept copy changed or is gone"],
)
class TestCleanupJob(DuplicateFilesTestCase):
async def test_cleanup_is_a_records_job_with_a_safety_backup(self):
self.trio()
plan, cleanup = await self.remove_all_but("a.png")
summary = cleanup.summary()
self.assertEqual(summary["kind"], "duplicate_cleanup")
self.assertEqual(summary["title"], "Remove duplicate image files")
self.assertEqual(summary["admission"], "records")
self.assertFalse(summary["cancellable"])
self.assertEqual(summary["scope"], {"operation": "duplicate_files"})
self.assertEqual(cleanup.state, "finished", cleanup.log)
self.assertEqual(cleanup.counters, {"removed": 2, "kept": 1, "errors": 0})
self.assertEqual(
sorted(row["path"] for row in cleanup.outcomes["removed"]),
["b.png", "sub/c.png"],
)
self.assertEqual(cleanup.outcomes["kept"][0]["path"], "a.png")
[backup] = self.backups()
self.assertEqual(
cleanup.status_line,
f"Removed 2 duplicate files · Safety backup {backup}",
)
self.assertEqual(
_sha(self.output_dir / "a.png"), hashlib.sha256(b"same bytes").hexdigest()
)
self.assertFalse((self.output_dir / "b.png").exists())
self.assertFalse((self.output_dir / "sub" / "c.png").exists())
self.assertNotIn(self.tmpdir, json.dumps(cleanup.report()))
# Single use, like every plan.
again = await self.apply(plan["token"], 409)
self.assertEqual(again["code"], "plan-used")
async def test_thumbnails_of_removed_files_are_removed(self):
self.trio()
self.write("clip.mp4", b"video")
self.write("copy.mp4", b"video")
kept = self.write("thumbnails/a_thumb.png", b"t")
gone = self.write("thumbnails/b_thumb.png", b"t")
video_thumb = self.write("thumbnails/copy_thumb.jpg", b"t")
job = await self.find()
groups = (await self.groups(job.id))["groups"]
keep = []
for group in groups:
paths = sorted(f["path"] for f in group["files"])
keep.append(
self.keep_of(group, "a.png" if "a.png" in paths else "clip.mp4")
)
plan = await self.review(job.id, keep)
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
self.assertTrue(kept.is_file())
self.assertFalse(gone.exists())
self.assertFalse(video_thumb.exists())
async def test_changed_content_skips_file(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.write("b.png", b"new bytes") # changed since the review
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
self.assertEqual(cleanup.counters, {"removed": 1, "kept": 2, "errors": 0})
self.assertEqual((self.output_dir / "b.png").read_bytes(), b"new bytes")
self.assertFalse((self.output_dir / "sub" / "c.png").exists())
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
self.assertEqual(rows["b.png"], duplicate_files.CHANGED)
async def test_a_changed_keeper_keeps_its_whole_group(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
self.write("a.png", b"edited keeper")
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
self.assertEqual(cleanup.counters["removed"], 0)
self.assertTrue((self.output_dir / "b.png").is_file())
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
self.assertTrue(
cleanup.status_line.startswith(
"Removed 0 duplicate files, 2 not removed · "
)
)
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
self.assertEqual(rows["b.png"], duplicate_files.KEEPER_CHANGED)
async def test_partial_failure_reports_each_path(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
real = duplicate_files._unlink
def unlink(path):
if os.path.basename(path) == "c.png":
raise PermissionError(13, f"Permission denied: {path}")
real(path)
self.write("z.png", b"same bytes") # never reviewed: not touched
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
self.assertEqual(cleanup.state, "failed")
self.assertEqual(cleanup.counters, {"removed": 1, "kept": 1, "errors": 1})
self.assertEqual(
[row["path"] for row in cleanup.outcomes["removed"]], ["b.png"]
)
[error] = cleanup.outcomes["errors"]
self.assertEqual(
(error["path"], error["detail"]),
("sub/c.png", duplicate_files.LOCKED_FILE),
)
self.assertEqual(
cleanup.status_line,
"Stopped after an error · Removed 1 duplicate file, 1 not removed · "
"Safety backup " + self.backups()[0],
)
report = json.dumps(cleanup.report())
self.assertNotIn("Permission denied", report)
self.assertNotIn("Errno", report)
self.assertTrue((self.output_dir / "z.png").is_file())
async def test_stop_on_failure_leaves_later_files_untouched(self):
self.write("a.png", b"same", mtime=1)
for name in ("b.png", "c.png", "e.png"):
self.write(name, b"same", mtime=2)
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
with (
patch.object(
duplicate_files, "_unlink", side_effect=OSError(5, "I/O error")
),
self.assertLogs("prompt_manager", "WARNING"),
):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
self.assertEqual(cleanup.counters, {"removed": 0, "kept": 3, "errors": 1})
self.assertEqual(
cleanup.outcomes["errors"][0]["detail"], duplicate_files.DELETE_FAILED
)
not_attempted = [
row
for row in cleanup.outcomes["kept"]
if row["detail"] == duplicate_files.NOT_ATTEMPTED
]
self.assertEqual(len(not_attempted), 2)
for name in ("b.png", "c.png", "e.png"):
self.assertTrue((self.output_dir / name).is_file())
async def test_locked_file_gives_a_coded_message(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
target = self.output_dir / "b.png"
handle = None
if os.name == "nt": # pragma: no cover - the native sharing violation
handle = open(target, "rb")
context = patch.object(duplicate_files, "_unlink", duplicate_files._unlink)
else:
context = patch.object(
duplicate_files,
"_unlink",
side_effect=PermissionError(
32, "The process cannot access the file", str(target)
),
)
try:
with context:
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
finally:
if handle is not None: # pragma: no cover - Windows
handle.close()
[error] = cleanup.outcomes["errors"]
self.assertEqual(error["detail"], duplicate_files.LOCKED_FILE)
self.assertNotIn("WinError", json.dumps(cleanup.report()))
self.assertTrue(target.is_file())
async def test_handles_closed_before_unlink(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
real = duplicate_files._unlink
open_at_unlink = []
def unlink(path):
open_at_unlink.extend(open_handles_under(path))
real(path)
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
self.assertEqual(cleanup.counters["removed"], 2)
self.assertEqual(open_at_unlink, [])
async def test_db_links_reconciled_only_for_removed_files(self):
self.trio()
db = self.api.db
moved = self.save("Moves to the keeper")
merged = self.save("Already on the keeper")
db.link_image_to_prompt(moved, str(self.output_dir / "b.png"))
db.link_image_to_prompt(moved, str(self.output_dir / "d.png"))
db.link_image_to_prompt(merged, str(self.output_dir / "a.png"))
db.link_image_to_prompt(
merged,
str(self.output_dir / "sub" / "c.png"),
{"workflow": {"only": "c"}, "prompt": {"only": "c"}},
)
db.set_image_favorite(
duplicate_files.favorite_root_key(self.output_dir), "b.png"
)
plan, cleanup = await self.remove_all_but("a.png")
# The merged link's own workflow document went with it (#249 S6 dedupe).
documents = db.model.get_connection().execute(
"SELECT COUNT(*) FROM workflow_documents"
)
self.assertEqual(documents.fetchone()[0], 1) # the shared "{}"
self.assertEqual(
plan["details"][0]["detail"],
"output · a copy of a.png · 1 database link moves to the kept copy",
)
def files_of(prompt_id):
return sorted(
os.path.relpath(image["image_path"], self.output_dir)
for image in db.get_prompt_images(prompt_id)
)
self.assertEqual(files_of(moved), ["a.png", "d.png"])
self.assertEqual(files_of(merged), ["a.png"])
root_key = duplicate_files.favorite_root_key(self.output_dir)
self.assertEqual(db.get_root_favorites(root_key), {"a.png"})
self.assertIn(
"Moved 2 database links to the kept copies", "\n".join(cleanup.log)
)
async def test_a_failed_removal_leaves_its_links_on_the_kept_copy(self):
self.trio()
prompt = self.save("Stays on c")
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "sub" / "c.png"))
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
real = duplicate_files._unlink
def unlink(path):
if os.path.basename(path) == "c.png":
raise PermissionError(13, "denied")
real(path)
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
# Its records moved to the identical kept copy before the unlink failed.
[image] = self.api.db.get_prompt_images(prompt)
self.assertTrue(image["image_path"].endswith("a.png"))
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
async def test_cleanup_and_find_exclude_each_other(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
gate, _calls = self.gate_hashing()
find = (await (await self.client.request("POST", FIND)).json())["job"]
busy = await self.apply(plan["token"], 409)
self.assertEqual((busy["code"], busy["running"]["id"]), ("busy", find["id"]))
gate.set()
await asyncio.wait_for(
asyncio.shield(self.api._jobs.get(find["id"]).task), DISK_BOUND_SECONDS
)
# The token survived the refusal; the cleanup holds off a new find.
release = asyncio.Event()
real = self.api._run_duplicate_cleanup
async def held(*args):
await release.wait()
return await real(*args)
with patch.object(self.api, "_run_duplicate_cleanup", side_effect=held):
body = await self.apply(plan["token"])
refused = await self.client.request("POST", FIND)
self.assertEqual(refused.status, 409)
legacy = await self.client.request("POST", STREAM)
self.assertEqual(legacy.status, 409)
release.set()
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
async def test_cleanup_refused_up_front_when_the_backup_does_not_fit(self):
from py.api import maintenance
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
with patch.object(maintenance, "_free_bytes", return_value=0):
reply = await self.apply(plan["token"], 507)
self.assertEqual(reply["code"], "insufficient-space")
self.assertTrue((self.output_dir / "b.png").is_file())
async def test_policy_off_removes_without_a_backup(self):
PromptManagerConfig.SAFETY_BACKUP = False
self.trio()
_plan, cleanup = await self.remove_all_but("a.png")
self.assertEqual(self.backups(), [])
self.assertEqual(
cleanup.status_line, "Removed 2 duplicate files · No safety backup"
)
class TestLegacyDuplicateRoutes(DuplicateFilesTestCase):
async def test_legacy_delete_duplicates_unchanged(self):
self.trio()
response = await self.client.request(
"POST", LEGACY_DELETE, json={"image_paths": ["b.png", "nope.png"]}
)
body = await response.json()
self.assertEqual(response.status, 200)
self.assertEqual(
set(body),
{"success", "deleted_count", "failed_count", "message", "failed_files"},
)
self.assertEqual((body["deleted_count"], body["failed_count"]), (1, 1))
self.assertEqual(body["failed_files"], ["nope.png (file not found)"])
[hold] = self.api._jobs.list()
self.assertEqual(
(hold.kind, hold.title, hold.admission),
("duplicate_cleanup", "Duplicate image delete (gallery page)", "records"),
)
async def test_legacy_delete_refused_while_a_find_runs(self):
self.trio()
gate, _calls = self.gate_hashing()
await self.client.request("POST", FIND)
response = await self.client.request(
"POST", LEGACY_DELETE, json={"image_paths": ["b.png"]}
)
self.assertEqual(response.status, 409)
self.assertEqual(
await response.json(),
{"success": False, "error": "Find duplicate image files already running"},
)
self.assertTrue((self.output_dir / "b.png").is_file())
gate.set()
async def test_legacy_scan_holds_the_duplicates_kind(self):
self.trio()
gate, _calls = self.gate_hashing()
legacy = asyncio.ensure_future(
self.client.request("GET", "/prompt_manager/scan_duplicates")
)
await self.until(lambda: self.api._jobs.active("duplicates") is not None)
[hold] = self.api._jobs.list()
self.assertEqual(
(hold.kind, hold.title, hold.admission, hold.cancellable),
("duplicates", "Find duplicate image files (gallery page)", "read", False),
)
refused = await self.client.request("POST", FIND)
self.assertEqual(refused.status, 409)
# A held legacy scan publishes no events: the stream does not attach.
stream = await self.client.request("POST", STREAM)
self.assertEqual(stream.status, 409)
gate.set()
body = await (await legacy).json()
self.assertEqual(set(body), {"success", "duplicates", "message"})
self.assertEqual(body["duplicates"][0]["count"], 3)
async def test_legacy_stream_does_not_attach_to_a_find_missing_its_folder(self):
other = self.second_source()
self.write("x.png", b"x", folder=other)
GalleryConfig.MONITORING_DIRECTORIES = [self._root(other)]
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
gate, _calls = self.gate_hashing()
await self.client.request("POST", FIND)
stream = await self.client.request("POST", STREAM)
self.assertEqual(stream.status, 409)
gate.set()
class TestLinkQueries(DuplicateFilesTestCase):
def test_link_queries_use_the_path_indexes(self):
conn = sqlite3.connect(self.api.db.model.db_path)
try:
for sql in duplicate_files.query_plan_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 generated_images", plan)
self.assertNotIn("SCAN image_favorites", plan)
finally:
conn.close()
class TestReviewHelpers(DuplicateFilesTestCase):
"""The pure checks of py/api/duplicate_review.py."""
def test_file_problem_on_disk(self):
from py.api import duplicate_review as review
root = str(self.output_dir)
good = self.write("ok.png", b"bytes")
digest = hashlib.sha256(b"bytes").hexdigest()
self.assertIsNone(review.file_problem(str(good), root, digest, 5))
self.assertEqual(
review.file_problem(str(good), root, digest, 6), review.CHANGED_CONTENT
)
self.assertEqual(
review.file_problem(str(good), root, "0" * 64, 5), review.CHANGED_CONTENT
)
self.assertEqual(
review.file_problem(os.path.join(root, "gone.png"), root, digest, 5),
review.GONE,
)
folder = self.output_dir / "folder.png"
folder.mkdir()
self.assertEqual(review.file_problem(str(folder), root, digest, 5), review.LINK)
with patch.object(review.os, "lstat", side_effect=OSError("io")):
self.assertEqual(
review.file_problem(str(good), root, digest, 5), review.UNREADABLE
)
def unreadable(_path):
raise OSError("io")
self.assertEqual(
review.file_problem(str(good), root, digest, 5, unreadable),
review.UNREADABLE,
)
def vanished(_path):
raise FileNotFoundError
self.assertEqual(
review.file_problem(str(good), root, digest, 5, vanished), review.GONE
)
def test_a_linked_folder_is_outside_its_source(self):
from py.api import duplicate_review as review
if not TestReview._can_symlink(self.tmpdir): # pragma: no cover
self.skipTest("symlinks are not available")
target = self.write("x.png", b"bytes", folder=self.outside_dir)
os.symlink(str(self.outside_dir), str(self.output_dir / "linked"))
digest = hashlib.sha256(b"bytes").hexdigest()
path = os.path.join(str(self.output_dir), "linked", "x.png")
self.assertTrue(os.path.isfile(path))
self.assertEqual(
review.file_problem(path, str(self.output_dir), digest, 5),
review.OUTSIDE_SOURCE,
)
self.assertTrue(target.is_file())
def test_options_are_validated(self):
from py.api import duplicate_review as review
keep = {"group_id": "g", "root_index": 0, "path": "a.png"}
cases = [
None,
{"job_id": "x"},
{"job_id": "", "keep": [keep]},
{"job_id": "x" * 65, "keep": [keep]},
{"job_id": "x", "keep": {}},
{"job_id": "x", "keep": [keep] * (review.MAX_KEEP + 1)},
{"job_id": "x", "keep": [{**keep, "group_id": 1}]},
{"job_id": "x", "keep": [{**keep, "root_index": True}]},
{"job_id": "x", "keep": [{**keep, "root_index": -1}]},
{"job_id": "x", "keep": [{**keep, "path": ""}]},
{"job_id": "x", "keep": [["g", 0, "a.png"]]},
]
for options in cases:
with self.subTest(options=options):
self.assertIsNotNone(review._parse_options(options)[2])
job_id, parsed, message = review._parse_options(
{"job_id": "x", "keep": [{**keep, "root_index": None}]}
)
self.assertEqual((job_id, parsed, message), ("x", [("g", None, "a.png")], None))
def test_sizes_read_naturally(self):
from py.api.duplicate_review import human_bytes
self.assertEqual(human_bytes(10), "10 B")
self.assertEqual(human_bytes(1536), "1.5 KB")
self.assertEqual(human_bytes(3 * 2**20), "3.0 MB")
self.assertEqual(human_bytes(2**40 * 5000), "5,000.0 TB")
async def test_a_kept_copy_listed_twice_counts_once(self):
self.trio()
job, group = await self.found_group()
keep = self.keep_of(group, "a.png")
plan = await self.review(job.id, [keep, keep])
self.assertEqual(plan["scope"]["count"], 2)
async def test_preview_failure_is_coded(self):
self.trio()
job, group = await self.found_group()
with (
patch.object(
self.api, "_review_duplicate_files", side_effect=OSError("boom")
),
self.assertLogs("prompt_manager", "ERROR"),
):
reply = await self.review(job.id, [self.keep_of(group, "a.png")], 500)
self.assertEqual(reply["error"], "Failed to review maintenance")
async def test_a_source_removed_before_the_apply_deletes_nothing(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.second_source())]
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
self.assertEqual(cleanup.state, "failed")
self.assertEqual(cleanup.last_event["code"], "sources-changed")
self.assertEqual(self.backups(), [])
self.assertTrue((self.output_dir / "b.png").is_file())
async def test_a_records_move_failure_deletes_nothing(self):
self.trio()
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
with (
patch.object(
self.api.db,
"move_copy_records",
side_effect=sqlite3.OperationalError("x"),
),
self.assertLogs("prompt_manager", "ERROR"),
):
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), DISK_BOUND_SECONDS)
self.assertEqual(cleanup.state, "failed")
self.assertEqual(
cleanup.outcomes["errors"][-1]["detail"], duplicate_files.RECORDS_FAILED
)
self.assertEqual(cleanup.counters["removed"], 0)
self.assertTrue((self.output_dir / "b.png").is_file())
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
Reference in New Issue View Git Blame Copy Permalink