Files
ComfyAssets-ComfyUI_PromptM…/tests/test_thumbnail_convert.py
T
Vito Sansevero 80da6f125e feat(thumbnails): Convert outdated thumbnails as an Activity job
POST /jobs/thumbnails {mode: "convert"} starts "Convert outdated
thumbnails" under the thumbnail job's kind, guarded admission, worker
threads and cancel at a file boundary. For each file whose preview is
in another format, it writes the new preview from the source (atomic,
verified; a video with no decodable frame fails instead of getting the
placeholder), then removes that file's old-format previews in its
thumbnails folder. A failure keeps the old preview; orphan previews
and links are never touched. The final line reports converted, failed
and bytes before and after. The plan counts it with mode=convert; the
legacy gallery stream refuses to attach to a running conversion.
2026-10-24 13:25:57 -07:00

282 lines
12 KiB
Python

"""Convert outdated thumbnails (#317): an explicit, cancellable Activity job.
``POST /prompt_manager/jobs/thumbnails {mode: "convert"}`` runs under the
thumbnail job's kind, admission (``guarded``) and worker threads. For every
media file whose preview is in another format than the one written now, it
writes the new preview from the SOURCE (verified, atomic), then removes that
file's old-format previews under ``thumbnails/``. A failure keeps the old
preview; previews whose source is gone (orphans) are never touched. The final
line reports converted / failed and bytes before → after.
Ordering is driven by events, never by timers.
"""
import asyncio
import os
import sys
import threading
import unittest
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
try:
from tests.test_api_images import ThumbnailTestCase, make_png
except ImportError: # discovered with tests/ as the top-level directory
from test_api_images import ThumbnailTestCase, make_png
from PIL import Image # noqa: E402
from py.api import thumbnail_encode # noqa: E402
from py.config import PromptManagerConfig # noqa: E402
JOB = "/prompt_manager/jobs/thumbnails"
PLAN = "/prompt_manager/images/thumbnails/plan"
STATUS = "/prompt_manager/images/thumbnails/status"
WAIT = 10
def big_png(path, size=(512, 512)):
"""A noisy PNG: its PNG preview is large, a WebP one much smaller."""
path.parent.mkdir(parents=True, exist_ok=True)
Image.effect_noise(size, 40).convert("RGB").save(path, "PNG")
return path
class ConvertCase(ThumbnailTestCase):
async def asyncSetUp(self):
await super().asyncSetUp()
saved = PromptManagerConfig.get_config()["thumbnails"]
self.addCleanup(PromptManagerConfig.update_config, {"thumbnails": saved})
PromptManagerConfig.update_config(
{"thumbnails": {"format": "webp", "quality": 80, "edge": 512}}
)
thumbnail_encode.reset_encoders()
self.addCleanup(thumbnail_encode.reset_encoders)
self.addCleanup(self.api.close_auto_thumbnails)
def _legacy(self, name, suffix=".png", source=True):
"""A media file (unless ``source`` is False) and its old-format preview."""
if source:
big_png(self.output_dir / f"{name}.png")
preview = big_png(self._thumb_path(name, suffix), size=(256, 256))
if source:
later = (self.output_dir / f"{name}.png").stat().st_mtime + 5
os.utime(preview, (later, later))
return preview
async def _convert(self):
resp = await self.client.post(JOB, json={"mode": "convert"})
body = await resp.json()
self.assertEqual(resp.status, 202, body)
job = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(job.task, WAIT)
return body["job"], job
class TestConvertOutdated(ConvertCase):
async def test_plan_counts_the_outdated_previews(self):
self._legacy("a")
self._legacy("b")
make_png(self.output_dir / "new.png") # missing: not this job's work
plan = await (await self.client.get(PLAN, params={"mode": "convert"})).json()
self.assertEqual(
(plan["mode"], plan["targets"], plan["outdated"]), ("convert", 2, 2)
)
self.assertEqual(self.api._jobs.list(), [])
async def test_converts_from_the_source_and_removes_the_old_preview(self):
old = self._legacy("a")
original = (self.output_dir / "a.png").read_bytes()
orphan = self._legacy("gone", source=False)
make_png(self.output_dir / "new.png")
before = old.stat().st_size
summary, job = await self._convert()
self.assertEqual(
(summary["kind"], summary["admission"], summary["cancellable"]),
("thumbnails", "guarded", True),
)
self.assertEqual(summary["title"], "Convert outdated thumbnails")
self.assertEqual(summary["scope"]["mode"], "convert")
self.assertEqual(job.state, "finished")
self.assertEqual(job.counters["generated"], 1)
self.assertEqual(job.counters["errors"], 0)
new = self._thumb_path("a")
with Image.open(new) as img:
self.assertEqual((img.format, max(img.size)), ("WEBP", 512))
self.assertFalse(old.exists(), "the old PNG preview is removed")
self.assertEqual((self.output_dir / "a.png").read_bytes(), original)
self.assertTrue(orphan.exists(), "an orphan preview is never touched")
self.assertFalse(self._thumb_path("new").exists(), "missing is not converted")
after = new.stat().st_size
self.assertRegex(job.status_line, r"^1 converted · 0 failed · .+ → .+$")
self.assertTrue(any("Conversion finished" in line for line in job.log))
self.assertLess(after, before)
status = await (await self.client.get(STATUS)).json()
self.assertEqual((status["outdated"], status["pending"]), (0, 1))
async def test_a_source_that_fails_keeps_its_old_preview(self):
old = self._legacy("bad")
(self.output_dir / "bad.png").write_bytes(b"\x89PNG broken")
later = old.stat().st_mtime + 10
os.utime(old, (later, later))
keep = old.read_bytes()
_, job = await self._convert()
self.assertEqual(job.counters["errors"], 1)
self.assertEqual(job.counters["generated"], 0)
self.assertEqual(old.read_bytes(), keep)
self.assertFalse(self._thumb_path("bad").exists())
self.assertRegex(job.status_line, r"^0 converted · 1 failed")
async def test_a_video_without_a_frame_keeps_its_old_preview(self):
(self.output_dir / "clip.mp4").write_bytes(b"no decoder")
old = big_png(self._thumb_path("clip", ".jpg"), size=(64, 64))
keep = old.read_bytes()
with (
patch.object(self.api, "_opencv_frame", return_value=None),
patch.object(self.api, "_ffmpeg_frame", return_value=None),
):
_, job = await self._convert()
self.assertEqual(job.counters["errors"], 1)
self.assertEqual(old.read_bytes(), keep)
self.assertFalse(self._thumb_path("clip").exists(), "never the placeholder")
async def test_an_old_preview_that_is_a_link_is_left(self):
big_png(self.output_dir / "l.png")
target = big_png(self.comfy_dir / "elsewhere.png", size=(32, 32))
link = self._thumb_path("l", ".png")
link.parent.mkdir(parents=True, exist_ok=True)
self._make_symlink(target, link)
_, job = await self._convert()
self.assertTrue(link.is_symlink())
self.assertTrue(target.exists())
self.assertEqual(job.counters["errors"], 0)
async def test_a_preview_already_current_only_loses_its_leftovers(self):
old = self._legacy("c")
current = big_png(self._thumb_path("c"), size=(64, 64))
later = old.stat().st_mtime + 5
os.utime(current, (later, later))
written = current.read_bytes()
_, job = await self._convert()
self.assertEqual(job.counters["generated"], 1)
self.assertEqual(current.read_bytes(), written, "not written again")
self.assertFalse(old.exists())
async def test_cancel_stops_at_a_file_boundary_keeping_what_converted(self):
original = PromptManagerConfig.WORKER_THREADS
PromptManagerConfig.WORKER_THREADS = 1
self.addCleanup(setattr, PromptManagerConfig, "WORKER_THREADS", original)
olds = [self._legacy(f"p{i}") for i in range(6)]
done = []
real = self.api._convert_one
job_ids = []
def counting(*args, **kwargs):
result = real(*args, **kwargs)
done.append(result)
if len(done) == 2:
self.assertTrue(self.api._jobs.cancel(job_ids[0]))
return result
self.api._convert_one = counting
resp = await self.client.post(JOB, json={"mode": "convert"})
job_ids.append((await resp.json())["job"]["id"])
job = self.api._jobs.get(job_ids[0])
await asyncio.wait_for(job.task, WAIT)
self.assertEqual(job.state, "cancelled")
self.assertEqual(job.counters["generated"], 2)
converted = [p for p in olds if not p.exists()]
self.assertEqual(len(converted), 2)
for old in olds:
webp = old.with_suffix(".webp")
self.assertEqual(webp.exists(), not old.exists(), old.name)
self.assertTrue(any("Conversion cancelled" in line for line in job.log))
async def test_refused_beside_a_running_thumbnail_job_or_scan(self):
self._legacy("a")
entered, release = threading.Event(), threading.Event()
self.addCleanup(release.set)
real = self.api._convert_one
def gated(*args, **kwargs):
entered.set()
release.wait(WAIT)
return real(*args, **kwargs)
self.api._convert_one = gated
resp = await self.client.post(JOB, json={"mode": "convert"})
first = (await resp.json())["job"]["id"]
loop = asyncio.get_running_loop()
self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
again = await self.client.post(JOB, json={"mode": "missing"})
self.assertEqual(again.status, 409)
scan = await self.client.post("/prompt_manager/jobs/scan")
self.assertEqual(scan.status, 409)
release.set()
await asyncio.wait_for(self.api._jobs.get(first).task, WAIT)
async def test_the_legacy_gallery_stream_never_attaches_to_a_conversion(self):
try:
from tests.test_api_images import parse_sse
except ImportError:
from test_api_images import parse_sse
self._legacy("a")
entered, release = threading.Event(), threading.Event()
self.addCleanup(release.set)
real = self.api._convert_one
def gated(*args, **kwargs):
entered.set()
release.wait(WAIT)
return real(*args, **kwargs)
self.api._convert_one = gated
resp = await self.client.post(JOB, json={"mode": "convert"})
first = (await resp.json())["job"]["id"]
loop = asyncio.get_running_loop()
self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
stream = await self.client.post(self.PROGRESS_URL)
events = parse_sse(await stream.text())
self.assertEqual(events[-1][0], "error")
self.assertIn("Convert outdated thumbnails", events[-1][1]["error"])
release.set()
await asyncio.wait_for(self.api._jobs.get(first).task, WAIT)
async def test_nothing_outdated_says_so(self):
big_png(self.output_dir / "a.png")
_, job = await self._convert()
self.assertEqual(job.state, "finished")
self.assertEqual(job.status_line, "No outdated previews to convert")
async def test_the_source_vanishing_keeps_the_old_preview(self):
old = self._legacy("v")
real = self.api._generate_one
def vanish(src, *args, **kwargs):
os.remove(src)
return real(src, *args, **kwargs)
self.api._generate_one = vanish
_, job = await self._convert()
self.assertEqual(job.counters["errors"], 1)
self.assertTrue(old.exists(), "a source gone mid-run makes an orphan")
def test_byte_sizes_are_readable(self):
from py.api.thumbnail_convert import readable_bytes
self.assertEqual(readable_bytes(0), "0 B")
self.assertEqual(readable_bytes(1536), "1.5 KB")
self.assertEqual(readable_bytes(5 * 1024 * 1024), "5.0 MB")
self.assertEqual(readable_bytes(3 * 1024**3), "3.0 GB")
if __name__ == "__main__":
unittest.main()