Vito Sansevero
38f5491da1
A planned target keeps the plan's own containment rule (inside the store, as the thumbnail job writes it); captures keep the serve rule. The structural test also covers PREVIEW_FOLDER joins and the repo-root modules; the write-lock test asserts the lock instead of waiting on a timer; new tests cover nested roots in the after-scan queue and a legacy upper-case preview served through the resolver.
458 lines
18 KiB
Python
458 lines
18 KiB
Python
"""One resolver for where previews live on disk (#320 PR 1).
|
|
|
|
``py/api/preview_store.py`` owns every mapping between an image source, a
|
|
media file and its preview on disk. In source mode (the only mode so far)
|
|
it returns exactly the paths the code built by hand before
|
|
(``<source>/thumbnails/<rel>_thumb<ext>``), so nothing changes on disk. The
|
|
structural test keeps new ``thumbnails`` path joins out of the rest of the
|
|
backend.
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import re
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from pathlib import Path
|
|
import unittest.mock
|
|
from unittest.mock import MagicMock
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
sys.modules.setdefault("server", MagicMock())
|
|
|
|
from py.api import preview_store, thumbnail_settings # noqa: E402
|
|
from py.api.thumbnail_encode import ( # noqa: E402
|
|
PreviewWriteHeld,
|
|
replace_unless_held,
|
|
)
|
|
|
|
REPO = Path(__file__).resolve().parent.parent
|
|
WAIT = 10
|
|
|
|
|
|
class SourceModeCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self._tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self._tmp.cleanup)
|
|
self.root = Path(self._tmp.name) / "output"
|
|
self.root.mkdir()
|
|
|
|
|
|
class TestSourceModePaths(SourceModeCase):
|
|
def test_default_location_is_inside_each_source(self):
|
|
location = preview_store.current()
|
|
self.assertEqual(location.mode, preview_store.SOURCE_MODE)
|
|
self.assertTrue(location.available)
|
|
self.assertIsNone(location.problem)
|
|
|
|
def test_store_of_a_source_is_its_thumbnails_folder(self):
|
|
self.assertEqual(preview_store.store_dir(self.root), self.root / "thumbnails")
|
|
|
|
def test_preview_path_matches_the_old_join(self):
|
|
rel = Path("sub", "deep", "a.png")
|
|
self.assertEqual(
|
|
preview_store.preview_path(self.root, rel, ".webp"),
|
|
self.root / "thumbnails/sub/deep/a_thumb.webp",
|
|
)
|
|
self.assertEqual(
|
|
preview_store.preview_path(self.root, "a.mp4", ".jpg"),
|
|
self.root / "thumbnails" / "a_thumb.jpg",
|
|
)
|
|
|
|
def test_preview_path_defaults_to_the_written_extension(self):
|
|
from py.api.thumbnail_encode import writer_settings
|
|
|
|
self.assertEqual(
|
|
preview_store.preview_path(self.root, "a.png").suffix,
|
|
writer_settings().extension,
|
|
)
|
|
|
|
def test_request_path_maps_url_parts_to_the_old_join(self):
|
|
parts = ("thumbnails", "sub", "a_thumb.webp")
|
|
self.assertEqual(
|
|
preview_store.request_path(self.root, parts), self.root.joinpath(*parts)
|
|
)
|
|
|
|
def test_request_path_is_none_for_anything_but_preview_parts(self):
|
|
self.assertIsNone(preview_store.request_path(self.root, ("sub", "a.png")))
|
|
self.assertIsNone(preview_store.request_path(self.root, ("Thumbnails", "a")))
|
|
self.assertIsNone(preview_store.request_path(self.root, ("thumbnails",)))
|
|
|
|
def test_ensure_store_is_the_store(self):
|
|
self.assertEqual(
|
|
preview_store.ensure_store(self.root), self.root / "thumbnails"
|
|
)
|
|
|
|
def test_owns_a_preview_inside_the_store(self):
|
|
inside = self.root / "thumbnails" / "sub" / "a_thumb.webp"
|
|
self.assertTrue(preview_store.owns(self.root, inside))
|
|
self.assertFalse(preview_store.owns(self.root, self.root / "a.png"))
|
|
self.assertFalse(
|
|
preview_store.owns(self.root, self.root.parent / "elsewhere.webp")
|
|
)
|
|
|
|
@unittest.skipUnless(hasattr(os, "symlink"), "needs symlinks")
|
|
def test_owns_refuses_a_store_linked_out_of_the_source(self):
|
|
outside = Path(self._tmp.name) / "outside"
|
|
outside.mkdir()
|
|
try:
|
|
(self.root / "thumbnails").symlink_to(outside, target_is_directory=True)
|
|
except OSError:
|
|
self.skipTest("symlinks not permitted")
|
|
self.assertFalse(
|
|
preview_store.owns(self.root, self.root / "thumbnails" / "a_thumb.webp")
|
|
)
|
|
|
|
@unittest.skipUnless(hasattr(os, "symlink"), "needs symlinks")
|
|
def test_owns_refuses_a_preview_folder_linked_out_of_the_store(self):
|
|
outside = Path(self._tmp.name) / "outside"
|
|
outside.mkdir()
|
|
(self.root / "thumbnails").mkdir()
|
|
try:
|
|
(self.root / "thumbnails" / "sub").symlink_to(
|
|
outside, target_is_directory=True
|
|
)
|
|
except OSError:
|
|
self.skipTest("symlinks not permitted")
|
|
self.assertFalse(
|
|
preview_store.owns(
|
|
self.root, self.root / "thumbnails" / "sub" / "a_thumb.webp"
|
|
)
|
|
)
|
|
|
|
def test_store_container_is_the_source(self):
|
|
self.assertEqual(preview_store.store_container(self.root), self.root)
|
|
|
|
def test_given_store_argument_is_kept_unless_it_is_the_default(self):
|
|
default = self.root / "thumbnails"
|
|
other = Path(self._tmp.name) / "custom"
|
|
self.assertEqual(preview_store.given_store(self.root, default), default)
|
|
self.assertEqual(preview_store.given_store(self.root, other), other)
|
|
|
|
|
|
class TestPreviewText(unittest.TestCase):
|
|
def test_preview_text_is_an_exact_thumbnails_prefix_and_preview_name(self):
|
|
for text in (
|
|
"thumbnails/a_thumb.webp",
|
|
"thumbnails/sub/a_thumb.jxl",
|
|
"thumbnails/a_thumb.JPG",
|
|
"thumbnails/a_thumb.Png",
|
|
):
|
|
with self.subTest(text=text):
|
|
self.assertTrue(preview_store.is_preview_text(text))
|
|
|
|
def test_other_text_is_not_preview_text(self):
|
|
for text in (
|
|
"Thumbnails/a_thumb.webp",
|
|
"thumbnails/a.png",
|
|
"thumbnails/_thumb.webp",
|
|
"thumbnails/a_thumb.txt",
|
|
"sub/thumbnails/a_thumb.webp",
|
|
"thumbnails/../a_thumb.webp",
|
|
"thumbnails//a_thumb.webp",
|
|
"thumbnails/./a_thumb.webp",
|
|
"thumbnails\\a_thumb.webp",
|
|
"thumbnails/c:a_thumb.webp",
|
|
"thumbnails/sub/thumbnails/a_thumb.webp",
|
|
"thumbnails",
|
|
"",
|
|
):
|
|
with self.subTest(text=text):
|
|
self.assertFalse(preview_store.is_preview_text(text))
|
|
|
|
|
|
class TestWritesAndTheSwitch(SourceModeCase):
|
|
def _temp_and_dst(self):
|
|
folder = self.root / "thumbnails"
|
|
folder.mkdir()
|
|
tmp = folder / ".a_thumb.part-1-1.webp"
|
|
tmp.write_bytes(b"new")
|
|
return tmp, folder / "a_thumb.webp"
|
|
|
|
def test_a_write_from_the_current_generation_lands(self):
|
|
tmp, dst = self._temp_and_dst()
|
|
replace_unless_held(tmp, dst, preview_store.write_guard())
|
|
self.assertEqual(dst.read_bytes(), b"new")
|
|
|
|
def test_a_write_whose_generation_changed_is_dropped(self):
|
|
tmp, dst = self._temp_and_dst()
|
|
guard = preview_store.write_guard()
|
|
preview_store.bump_generation()
|
|
with self.assertRaises(PreviewWriteHeld):
|
|
replace_unless_held(tmp, dst, guard)
|
|
self.assertFalse(dst.exists())
|
|
|
|
def test_the_hold_check_still_applies(self):
|
|
tmp, dst = self._temp_and_dst()
|
|
with self.assertRaises(PreviewWriteHeld):
|
|
replace_unless_held(tmp, dst, preview_store.write_guard(lambda: False))
|
|
self.assertFalse(dst.exists())
|
|
|
|
def test_the_check_and_the_rename_run_under_the_write_lock(self):
|
|
"""The hold check and the rename hold ``write_lock``, the lock the
|
|
generation bump takes: a switch cannot fall between them."""
|
|
tmp, dst = self._temp_and_dst()
|
|
in_check = threading.Event()
|
|
release = threading.Event()
|
|
seen = {}
|
|
|
|
def slow_hold():
|
|
seen["locked"] = preview_store.write_lock.locked()
|
|
in_check.set()
|
|
release.wait(WAIT)
|
|
return True
|
|
|
|
guard = preview_store.write_guard(slow_hold)
|
|
writer = threading.Thread(target=replace_unless_held, args=(tmp, dst, guard))
|
|
writer.start()
|
|
try:
|
|
self.assertTrue(in_check.wait(WAIT))
|
|
self.assertTrue(seen["locked"])
|
|
self.assertTrue(preview_store.write_lock.locked())
|
|
finally:
|
|
release.set()
|
|
writer.join(WAIT)
|
|
self.assertFalse(preview_store.write_lock.locked())
|
|
self.assertEqual(dst.read_bytes(), b"new")
|
|
bumped = preview_store.bump_generation()
|
|
self.assertEqual(preview_store.generation(), bumped)
|
|
|
|
|
|
class TestOnlyTheResolverBuildsStorePaths(unittest.TestCase):
|
|
"""No backend module joins ``thumbnails`` onto a path by hand any more."""
|
|
|
|
ALLOWED = {
|
|
Path("py/api/preview_store.py"),
|
|
Path("py/api/thumbnail_settings.py"), # URL text
|
|
}
|
|
# A path join with a literal "thumbnails" operand.
|
|
JOIN = re.compile(
|
|
r"""/\s*["']thumbnails["']"""
|
|
r"""|["']thumbnails["']\s*/"""
|
|
r"""|joinpath\([^)]*["']thumbnails["']"""
|
|
r"""|os\.path\.join\([^)]*["']thumbnails["']"""
|
|
r"""|["']thumbnails/"""
|
|
r"""|/\s*PREVIEW_FOLDER|joinpath\([^)]*PREVIEW_FOLDER"""
|
|
r"""|Path\([^)]*,\s*["']thumbnails["']"""
|
|
)
|
|
|
|
def test_no_thumbnails_join_outside_the_resolver(self):
|
|
offenders = []
|
|
files = sorted(REPO.glob("*.py"))
|
|
for folder in ("py", "utils", "database"):
|
|
files += sorted((REPO / folder).rglob("*.py"))
|
|
for path in files:
|
|
rel = path.relative_to(REPO)
|
|
if rel in self.ALLOWED:
|
|
continue
|
|
for number, line in enumerate(
|
|
path.read_text(encoding="utf-8").splitlines(), 1
|
|
):
|
|
code = line.split("#", 1)[0]
|
|
if self.JOIN.search(code):
|
|
offenders.append(f"{rel}:{number}: {line.strip()}")
|
|
self.assertEqual(offenders, [])
|
|
|
|
|
|
class TestSettingsLookupsFollowTheStore(SourceModeCase):
|
|
def test_owned_thumbnails_reads_the_store(self):
|
|
folder = self.root / "thumbnails" / "sub"
|
|
folder.mkdir(parents=True)
|
|
(folder / "a_thumb.webp").write_bytes(b"x")
|
|
self.assertEqual(
|
|
thumbnail_settings.owned_thumbnails(self.root, Path("sub", "a.png")),
|
|
[folder / "a_thumb.webp"],
|
|
)
|
|
|
|
|
|
try:
|
|
from tests.test_thumbnail_auto import AutoThumbnailCase, _entry
|
|
from tests.test_api_images import make_png
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from test_thumbnail_auto import AutoThumbnailCase, _entry
|
|
from test_api_images import make_png
|
|
|
|
|
|
class RedirectedStoreCase(AutoThumbnailCase):
|
|
"""Every caller follows the resolver: with the store of the output root
|
|
moved to ``<tmp>/elsewhere/output-key`` (a stand-in for a later mode),
|
|
previews are read, written, listed and removed there and only there."""
|
|
|
|
async def asyncSetUp(self):
|
|
await super().asyncSetUp()
|
|
self.container = self.comfy_dir / "elsewhere"
|
|
self.store = self.container / "output-key"
|
|
self.store.mkdir(parents=True)
|
|
self.old_store = self.output_dir / "thumbnails"
|
|
root = self.output_dir
|
|
real_store = preview_store.store_dir
|
|
real_container = preview_store.store_container
|
|
|
|
def store_dir(r, location=None):
|
|
return self.store if Path(r) == root else real_store(r, location)
|
|
|
|
def store_container(r, location=None):
|
|
if Path(r) == root:
|
|
return self.container
|
|
return real_container(r, location)
|
|
|
|
for name, fake in (
|
|
("store_dir", store_dir),
|
|
("store_container", store_container),
|
|
):
|
|
patcher = unittest.mock.patch.object(preview_store, name, fake)
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
def _preview(self, rel, ext=".webp"):
|
|
return self.store / (thumbnail_settings.preview_stem(rel) + ext)
|
|
|
|
|
|
class TestEveryCallerFollowsTheStore(RedirectedStoreCase):
|
|
async def test_an_existing_preview_is_served_from_the_store(self):
|
|
make_png(self.output_dir / "sub" / "a.png", size=(64, 64))
|
|
preview = self._preview("sub/a.png")
|
|
make_png(preview, size=(32, 32))
|
|
url = _entry(self.api, self.output_dir, "sub/a.png")["thumbnail_url"]
|
|
self.assertIn("v=", url) # the listing found it in the store
|
|
resp = await self.client.get(url)
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(await resp.read(), preview.read_bytes())
|
|
|
|
async def test_a_legacy_upper_case_preview_is_served_from_the_store(self):
|
|
make_png(self.output_dir / "k.JPG", size=(64, 64))
|
|
legacy = self.store / "k_thumb.JPG"
|
|
legacy.write_bytes(b"legacy preview")
|
|
resp = await self.client.get(
|
|
"/prompt_manager/images/serve/thumbnails/k_thumb.JPG?root=0"
|
|
)
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(await resp.read(), b"legacy preview")
|
|
|
|
async def test_a_missing_preview_is_made_in_the_store_on_first_view(self):
|
|
make_png(self.output_dir / "sub" / "b.png", size=(640, 480))
|
|
url = _entry(self.api, self.output_dir, "sub/b.png")["thumbnail_url"]
|
|
resp = await self.client.get(url)
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertTrue(self._preview("sub/b.png").is_file())
|
|
self.assertFalse(self.old_store.exists())
|
|
|
|
async def test_a_leftover_in_the_source_folder_is_not_served(self):
|
|
make_png(self.output_dir / "c.png", size=(640, 480))
|
|
leftover = self.old_store / "c_thumb.webp"
|
|
leftover.parent.mkdir()
|
|
leftover.write_bytes(b"old preview bytes")
|
|
resp = await self.client.get(
|
|
"/prompt_manager/images/serve/thumbnails/c_thumb.webp?root=0"
|
|
)
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertNotEqual(await resp.read(), b"old preview bytes")
|
|
self.assertTrue(self._preview("c.png").is_file())
|
|
|
|
async def test_capture_writes_into_the_store(self):
|
|
path = make_png(self.output_dir / "run" / "d.png", size=(64, 64))
|
|
self.assertEqual(self.api.queue_capture_thumbnail(str(path)), "queued")
|
|
await self._idle()
|
|
self.assertTrue(self._preview("run/d.png").is_file())
|
|
self.assertFalse(self.old_store.exists())
|
|
|
|
async def test_generate_and_clear_follow_the_store(self):
|
|
make_png(self.output_dir / "e.png", size=(64, 64))
|
|
payload = await asyncio.get_running_loop().run_in_executor(
|
|
None, self.api._generate_gallery_thumbnails_sync, [self.output_dir]
|
|
)
|
|
self.assertTrue(payload["success"])
|
|
self.assertTrue(self._preview("e.png").is_file())
|
|
self.assertFalse(self.old_store.exists())
|
|
self.assertTrue(payload["thumbnails_path"].endswith("output-key"))
|
|
cleared, _size = self.api._clear_gallery_thumbnails_sync([self.output_dir])
|
|
self.assertEqual(cleared, 1)
|
|
self.assertFalse(self._preview("e.png").exists())
|
|
self.assertTrue(self.store.is_dir()) # the store itself stays
|
|
|
|
async def test_removing_a_file_removes_its_preview_in_the_store(self):
|
|
from py.api.admin import _remove_thumbnail
|
|
|
|
make_png(self.output_dir / "f.png", size=(64, 64))
|
|
preview = self._preview("f.png")
|
|
preview.write_bytes(b"x")
|
|
self.assertTrue(_remove_thumbnail(self.output_dir, Path("f.png")))
|
|
self.assertFalse(preview.exists())
|
|
|
|
async def test_the_temp_sweep_walks_the_store(self):
|
|
stray = self.store / ".g_thumb.part-1-1.webp"
|
|
stray.write_bytes(b"x")
|
|
old = os.path.getmtime(stray) - 10 * 24 * 3600
|
|
os.utime(stray, (old, old))
|
|
self.assertEqual(self.api.sweep_thumbnail_temps(), 1)
|
|
self.assertFalse(stray.exists())
|
|
|
|
async def test_history_cover_writes_into_the_store(self):
|
|
from py.api import history_current
|
|
|
|
self.assertTrue(
|
|
history_current._preview_writable(self.output_dir, self._preview("h.png"))
|
|
)
|
|
self.assertFalse(
|
|
history_current._preview_writable(
|
|
self.output_dir, self.old_store / "h_thumb.webp"
|
|
)
|
|
)
|
|
|
|
|
|
class TestCaptureResolvesAtWriteTime(RedirectedStoreCase):
|
|
def test_a_queued_capture_follows_a_store_that_moved_before_its_write(self):
|
|
from py.api import thumbnail_auto as auto
|
|
|
|
written = []
|
|
thumbnailer = auto.AutoThumbnailer(
|
|
lambda src, dst, *a, **k: written.append(dst) or {"action": "generated"},
|
|
lambda: True,
|
|
self.api.logger,
|
|
start_worker=False,
|
|
)
|
|
self.addCleanup(thumbnailer.close)
|
|
src = make_png(self.output_dir / "i.png", size=(32, 32))
|
|
deferred = auto.DeferredPreview(self.output_dir, Path("i.png"))
|
|
self.assertTrue(thumbnailer.enqueue(src, deferred, False))
|
|
moved = self.container / "moved"
|
|
moved.mkdir()
|
|
self.store = moved # the location changes before the worker runs
|
|
thumbnailer._write_queued(thumbnailer._queue.popleft())
|
|
self.assertEqual(written, [moved / "i_thumb.webp"])
|
|
|
|
def test_the_after_scan_queue_maps_each_target_to_its_own_root(self):
|
|
from py.api import thumbnail_auto as auto
|
|
|
|
nested = self.output_dir / "inner"
|
|
nested.mkdir()
|
|
src = nested / "j.png"
|
|
outer = auto._deferred(
|
|
[self.output_dir, nested], src, self._preview("inner/j.png")
|
|
)
|
|
inner = auto._deferred(
|
|
[self.output_dir, nested], src, nested / "thumbnails" / "j_thumb.webp"
|
|
)
|
|
self.assertEqual(
|
|
(outer.root, outer.rel), (self.output_dir, Path("inner/j.png"))
|
|
)
|
|
self.assertEqual((inner.root, inner.rel), (nested, Path("j.png")))
|
|
stray = self.comfy_dir / "x_thumb.webp"
|
|
self.assertEqual(auto._deferred([self.output_dir], src, stray), stray)
|
|
|
|
def test_serve_and_history_share_one_flight_key(self):
|
|
from py.api import thumbnail_auto as auto
|
|
|
|
parts = ("thumbnails", "sub", "a_thumb.webp")
|
|
self.assertEqual(
|
|
auto._request_key(self.output_dir, parts),
|
|
auto._dst_key(self._preview("sub/a.png")),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|