ComfyAssets-ComfyUI_PromptManager/tests/test_api_admin.py

"""Route tests for the admin API (py/api/admin.py).

Uses aiohttp's test client against a PromptManagerAPI wired to a temporary
database and a stub ``folder_paths`` module, so the tests run without ComfyUI.
"""

import asyncio
import json
import os
import shutil
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import MagicMock

_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules.setdefault("server", _mock_server)

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.models import PromptModel  # noqa: E402

try:
    from tests.open_handles import assert_closed  # noqa: E402
except ImportError:  # discovered with tests/ as the top-level directory
    from open_handles import assert_closed  # noqa: E402

from aiohttp import web  # noqa: E402
from aiohttp.test_utils import AioHTTPTestCase  # noqa: E402

from database.operations import PromptDatabase  # noqa: E402
from py.api import PromptManagerAPI  # noqa: E402
from py.config import GalleryConfig, PromptManagerConfig  # noqa: E402

CONFIG_PATH_ENV = "PROMPT_MANAGER_CONFIG_PATH"
EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"


class AdminAPITestCase(AioHTTPTestCase):
    """Test app with PromptManager routes, temp DB and config, fake ComfyUI tree."""

    async def get_application(self):
        # realpath: macOS hands out /var/... for a /private/var/... directory and
        # the gallery roots are stored resolved
        self.tmpdir = os.path.realpath(tempfile.mkdtemp())
        self.comfy_dir = Path(self.tmpdir) / "ComfyUI"
        self.output_dir = self.comfy_dir / "output"
        self.output_dir.mkdir(parents=True)
        self.outside_dir = Path(self.tmpdir) / "outside"
        self.outside_dir.mkdir()

        self.config_path = os.path.join(self.tmpdir, "config.json")
        self._orig_env = {
            key: os.environ.pop(key, None) for key in (CONFIG_PATH_ENV, EXTRA_ROOTS_ENV)
        }
        os.environ[CONFIG_PATH_ENV] = self.config_path

        self._orig_folder_paths = sys.modules.get("folder_paths")
        sys.modules["folder_paths"] = types.SimpleNamespace(
            base_path=str(self.comfy_dir),
            get_output_directory=lambda: str(self.output_dir),
        )

        self._orig_gallery_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
        self._orig_pm_config = PromptManagerConfig.get_config()
        GalleryConfig.MONITORING_DIRECTORIES = []

        db_path = os.path.join(self.tmpdir, "prompts.db")
        app = web.Application()
        routes = web.RouteTableDef()
        self.api = PromptManagerAPI()
        self.api.db = PromptDatabase(db_path)
        self.api.add_routes(routes)
        app.router.add_routes(routes)
        return app

    async def tearDownAsync(self):
        await super().tearDownAsync()  # closes the aiohttp test client
        GalleryConfig.MONITORING_DIRECTORIES = self._orig_gallery_dirs
        PromptManagerConfig.update_config(self._orig_pm_config)
        for key, value in self._orig_env.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value
        if self._orig_folder_paths is None:
            sys.modules.pop("folder_paths", None)
        else:
            sys.modules["folder_paths"] = self._orig_folder_paths
        PromptModel.close_all_instances()  # restore may have swapped api.db
        assert_closed(self, self.tmpdir)
        shutil.rmtree(self.tmpdir, ignore_errors=True)

    # ── helpers ────────────────────────────────────────────────────────

    @staticmethod
    def _root(path):
        """A gallery root as the config stores it: resolved, case-folded on Windows."""
        return GalleryConfig.resolve_gallery_root(str(path))

    def _write_config(self, content):
        with open(self.config_path, "w") as f:
            json.dump(content, f)

    def _read_config(self):
        with open(self.config_path) as f:
            return json.load(f)


class TestSaveSettingsGalleryRoots(AdminAPITestCase):

    async def test_rejects_filesystem_root_and_keeps_config(self):
        self._write_config({"marker": "untouched"})

        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_paths": [os.path.abspath(os.sep)]},
        )

        self.assertEqual(resp.status, 400)
        data = await resp.json()
        self.assertFalse(data["success"])
        self.assertEqual(self._read_config(), {"marker": "untouched"})
        self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [])

    async def test_rejects_directory_outside_comfyui(self):
        self._write_config({"marker": "untouched"})

        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_paths": [self._root(self.outside_dir)]},
        )

        self.assertEqual(resp.status, 400)
        self.assertEqual(self._read_config(), {"marker": "untouched"})
        self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [])

    async def test_one_bad_root_rejects_whole_request(self):
        good = self.output_dir / "good"
        good.mkdir()
        self._write_config({"marker": "untouched"})

        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_paths": [str(good), str(self.outside_dir)]},
        )

        self.assertEqual(resp.status, 400)
        self.assertEqual(self._read_config(), {"marker": "untouched"})
        self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [])

    async def test_legacy_single_path_rejected_too(self):
        self._write_config({"marker": "untouched"})

        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_path": os.path.expanduser("~")},
        )

        self.assertEqual(resp.status, 400)
        self.assertEqual(self._read_config(), {"marker": "untouched"})

    async def test_accepts_subdir_of_output_and_persists(self):
        good = self.output_dir / "renders"
        good.mkdir()

        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_paths": [self._root(good)]},
        )

        self.assertEqual(resp.status, 200)
        data = await resp.json()
        self.assertTrue(data["success"])
        self.assertTrue(data["restart_required"])
        self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [self._root(good)])
        saved = self._read_config()
        self.assertEqual(
            saved["gallery"]["monitoring"]["directories"], [self._root(good)]
        )

    async def test_accepts_env_listed_directory(self):
        os.environ[EXTRA_ROOTS_ENV] = str(self.outside_dir)

        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_paths": [self._root(self.outside_dir)]},
        )

        self.assertEqual(resp.status, 200)
        self.assertEqual(
            GalleryConfig.MONITORING_DIRECTORIES, [self._root(self.outside_dir)]
        )

    async def test_legacy_single_path_accepted(self):
        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_path": str(self.output_dir)},
        )

        self.assertEqual(resp.status, 200)
        self.assertEqual(
            GalleryConfig.MONITORING_DIRECTORIES, [self._root(self.output_dir)]
        )

    async def test_empty_paths_clear_roots(self):
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_paths": ["", "   "]},
        )

        self.assertEqual(resp.status, 200)
        self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [])
        self.assertEqual(
            self._read_config()["gallery"]["monitoring"]["directories"], []
        )

    async def test_paths_must_be_a_list(self):
        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_paths": str(self.output_dir)},
        )
        self.assertEqual(resp.status, 400)

    async def test_path_entries_must_be_strings(self):
        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_paths": [42]},
        )
        self.assertEqual(resp.status, 400)


class FakeMonitor:
    def __init__(self, dirs):
        self.monitored_directories = list(dirs)

    def get_status(self):
        return {
            "running": True,
            "monitored_directories": self.monitored_directories,
            "handler_active": True,
            "observer_alive": True,
        }


class TestNoAbsolutePathsInResponses(AdminAPITestCase):
    """No admin response may reveal the server's absolute filesystem layout."""

    def _install_monitor(self, dirs):
        import utils.image_monitor as im_mod

        orig = im_mod._monitor_instance
        im_mod._monitor_instance = FakeMonitor(dirs)
        self.addCleanup(setattr, im_mod, "_monitor_instance", orig)

    async def _body(self, method, path, **kwargs):
        resp = await self.client.request(method, path, **kwargs)
        return resp, await resp.text()

    async def test_settings_show_roots_relative_to_comfyui(self):
        renders = self.output_dir / "renders"
        renders.mkdir()
        GalleryConfig.MONITORING_DIRECTORIES = [str(renders)]
        self._install_monitor([str(renders)])

        resp, body = await self._body("GET", "/prompt_manager/settings")

        self.assertEqual(resp.status, 200)
        self.assertNotIn(self.tmpdir, body)
        data = json.loads(body)
        self.assertEqual(data["settings"]["gallery_root_paths"], ["output/renders"])
        self.assertEqual(data["settings"]["gallery_root_path"], "output/renders")
        self.assertEqual(data["settings"]["monitored_directories"], ["output/renders"])

    async def test_settings_relative_root_round_trips(self):
        """What GET returns can be POSTed back unchanged from any CWD."""
        renders = self.output_dir / "renders"
        renders.mkdir()

        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_paths": ["output/renders"]},
        )

        self.assertEqual(resp.status, 200)
        self.assertEqual(
            [os.path.normcase(p) for p in GalleryConfig.MONITORING_DIRECTORIES],
            [os.path.normcase(os.path.realpath(renders))],
        )

    @unittest.skipIf(sys.platform == "win32", "symlinks need privileges on Windows")
    async def test_same_output_dir_spelled_twice_is_listed_once(self):
        alias = self.comfy_dir / "alias"
        os.symlink(self.output_dir, alias, target_is_directory=True)
        # Two spellings of one directory (Windows adds case differences too)
        GalleryConfig.MONITORING_DIRECTORIES = [str(alias), self._root(self.output_dir)]
        self._install_monitor([self._root(self.output_dir)])

        resp, body = await self._body("GET", "/prompt_manager/diagnostics")

        self.assertEqual(resp.status, 200)
        data = json.loads(body)
        self.assertEqual(
            data["diagnostics"]["comfyui_output"]["output_dirs"], ["output"]
        )
        self.assertEqual(len(self.api._get_all_output_dirs()), 1)

    async def test_diagnostics_paths_are_relative(self):
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        self._install_monitor([self._root(self.output_dir)])

        resp, body = await self._body("GET", "/prompt_manager/diagnostics")

        self.assertEqual(resp.status, 200)
        self.assertNotIn(self.tmpdir, body)
        data = json.loads(body)
        self.assertEqual(
            data["diagnostics"]["comfyui_output"]["output_dirs"], ["output"]
        )
        self.assertEqual(
            data["diagnostics"]["image_monitor"]["monitored_directories"], ["output"]
        )

    async def test_diagnostics_missing_db_message_has_no_path(self):
        self.api.db.model.db_path = os.path.join(self.tmpdir, "missing", "prompts.db")

        resp, body = await self._body("GET", "/prompt_manager/diagnostics")

        self.assertNotIn(self.tmpdir, body)
        data = json.loads(body)
        self.assertEqual(data["diagnostics"]["database"]["status"], "error")
        self.assertIn("prompts.db", data["diagnostics"]["database"]["message"])

    async def test_scan_duplicates_paths_are_relative_to_output(self):
        (self.output_dir / "a.png").write_bytes(b"same-bytes")
        (self.output_dir / "sub").mkdir()
        (self.output_dir / "sub" / "b.png").write_bytes(b"same-bytes")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

        resp, body = await self._body("GET", "/prompt_manager/scan_duplicates")

        self.assertEqual(resp.status, 200)
        self.assertNotIn(self.tmpdir, body)
        data = json.loads(body)
        self.assertEqual(len(data["duplicates"]), 1)
        paths = sorted(img["path"] for img in data["duplicates"][0]["images"])
        self.assertEqual(paths, ["a.png", os.path.join("sub", "b.png")])

    async def test_delete_duplicates_accepts_relative_paths(self):
        target = self.output_dir / "sub" / "b.png"
        target.parent.mkdir()
        target.write_bytes(b"bytes")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

        resp, body = await self._body(
            "POST",
            "/prompt_manager/delete_duplicate_images",
            json={"image_paths": [os.path.join("sub", "b.png")]},
        )

        self.assertEqual(resp.status, 200)
        data = json.loads(body)
        self.assertEqual(data["deleted_count"], 1)
        self.assertFalse(target.exists())

    async def test_scan_images_stream_has_no_absolute_paths(self):
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

        resp, body = await self._body("POST", "/prompt_manager/scan")

        self.assertEqual(resp.status, 200)
        self.assertNotIn(self.tmpdir, body)
        self.assertIn('"directories": ["output"]', body)


def _write_png(path, prompt_text=None):
    """Write a tiny PNG, optionally carrying a ComfyUI-style 'prompt' chunk."""
    from PIL import Image, PngImagePlugin

    img = Image.new("RGB", (2, 2), (255, 0, 0))
    info = PngImagePlugin.PngInfo()
    if prompt_text is not None:
        graph = {
            "1": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt_text}},
            "2": {"class_type": "CLIPTextEncode", "inputs": {"text": "bad"}},
            "3": {
                "class_type": "KSampler",
                "inputs": {"positive": ["1", 0], "negative": ["2", 0]},
            },
        }
        info.add_text("prompt", json.dumps(graph))
    img.save(str(path), pnginfo=info)


class TestScansRunOffTheEventLoop(AdminAPITestCase):
    """Blocking filesystem work must go through _run_in_executor."""

    def _spy_executor(self):
        calls = []
        original = self.api._run_in_executor

        async def spy(func, *args, **kwargs):
            calls.append(getattr(func, "__name__", repr(func)))
            return await original(func, *args, **kwargs)

        self.api._run_in_executor = spy
        return calls

    async def test_scan_duplicates_runs_in_executor(self):
        (self.output_dir / "a.png").write_bytes(b"dup")
        (self.output_dir / "b.png").write_bytes(b"dup")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        calls = self._spy_executor()

        resp = await self.client.request("GET", "/prompt_manager/scan_duplicates")

        self.assertEqual(resp.status, 200)
        self.assertEqual((await resp.json())["duplicates"][0]["count"], 2)
        self.assertIn("_find_duplicate_images_sync", calls)

    async def test_delete_duplicates_runs_in_executor(self):
        (self.output_dir / "a.png").write_bytes(b"dup")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        calls = self._spy_executor()

        resp = await self.client.request(
            "POST",
            "/prompt_manager/delete_duplicate_images",
            json={"image_paths": ["a.png"]},
        )

        self.assertEqual(resp.status, 200)
        self.assertEqual((await resp.json())["deleted_count"], 1)
        self.assertIn("_delete_duplicate_images_sync", calls)

    async def test_output_scan_runs_in_executor(self):
        _write_png(self.output_dir / "gen.png", prompt_text="a red square")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        calls = self._spy_executor()

        resp = await self.client.request("POST", "/prompt_manager/scan")

        body = await resp.text()
        self.assertEqual(resp.status, 200)
        self.assertIn('"type": "complete"', body)
        self.assertIn('"added": 1', body)
        for name in (
            "_collect_output_media_sync",
            "_extract_batch_metadata_sync",
            "_ingest_scan_batch_sync",
        ):
            self.assertIn(name, calls)

    async def test_output_scan_links_existing_prompt_on_second_pass(self):
        _write_png(self.output_dir / "gen.png", prompt_text="a red square")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

        first = await (await self.client.request("POST", "/prompt_manager/scan")).text()
        second = await (
            await self.client.request("POST", "/prompt_manager/scan")
        ).text()

        self.assertIn('"added": 1', first)
        self.assertIn('"linked": 1', second)
        self.assertIn('"added": 0', second)
        prompts = self.api.db.get_recent_prompts(limit=10)["prompts"]
        self.assertEqual([p["text"] for p in prompts], ["a red square"])

    async def test_output_scan_without_directories_reports_error(self):
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.outside_dir)]
        sys.modules.pop("folder_paths", None)
        self.api._find_comfyui_output_dir = lambda: None

        body = await (await self.client.request("POST", "/prompt_manager/scan")).text()

        self.assertIn('"type": "error"', body)


class TestDuplicateScanSyncBody(AdminAPITestCase):
    """Behaviour of the synchronous duplicate scan against a real tree."""

    def _scan(self):
        return self.api._find_duplicate_images_sync(str(self.output_dir))

    def test_groups_identical_files_and_sorts_by_mtime(self):
        older = self.output_dir / "older.png"
        newer = self.output_dir / "sub" / "newer.PNG"
        newer.parent.mkdir()
        older.write_bytes(b"same")
        newer.write_bytes(b"same")
        os.utime(older, (1_000_000, 1_000_000))
        os.utime(newer, (2_000_000, 2_000_000))
        (self.output_dir / "unique.png").write_bytes(b"different")

        groups = self._scan()

        self.assertEqual(len(groups), 1)
        self.assertEqual(groups[0]["count"], 2)
        self.assertEqual(
            [img["filename"] for img in groups[0]["images"]], ["older.png", "newer.PNG"]
        )
        self.assertEqual(
            groups[0]["images"][1]["path"], os.path.join("sub", "newer.PNG")
        )
        self.assertEqual(
            groups[0]["images"][1]["url"], "/prompt_manager/images/serve/sub/newer.PNG"
        )
        self.assertNotIn(self.tmpdir, json.dumps(groups))

    def test_thumbnails_dir_is_skipped_and_thumbnail_url_reported(self):
        (self.output_dir / "a.png").write_bytes(b"same")
        (self.output_dir / "b.png").write_bytes(b"same")
        thumbs = self.output_dir / "thumbnails"
        thumbs.mkdir()
        (thumbs / "a_thumb.png").write_bytes(b"same")
        (thumbs / "b_thumb.png").write_bytes(b"other")

        groups = self._scan()

        self.assertEqual(len(groups), 1)
        by_name = {img["filename"]: img for img in groups[0]["images"]}
        self.assertEqual(set(by_name), {"a.png", "b.png"})
        self.assertRegex(
            by_name["a.png"]["thumbnail_url"],
            r"^/prompt_manager/images/serve/thumbnails/a_thumb\.png\?v=\d+$",
        )
        self.assertRegex(
            by_name["b.png"]["thumbnail_url"],
            r"^/prompt_manager/images/serve/thumbnails/b_thumb\.png\?v=\d+$",
        )

    def test_videos_are_flagged_and_unreadable_entries_skipped(self):
        (self.output_dir / "clip.mp4").write_bytes(b"vid")
        (self.output_dir / "copy.mp4").write_bytes(b"vid")
        (self.output_dir / "not_a_file.png").mkdir()  # rglob matches, hashing fails

        groups = self._scan()

        self.assertEqual(len(groups), 1)
        self.assertTrue(all(img["is_video"] for img in groups[0]["images"]))
        self.assertTrue(
            all(img["media_type"] == "video" for img in groups[0]["images"])
        )

    def test_missing_output_dir_gives_empty_list(self):
        self.assertEqual(
            self.api._find_duplicate_images_sync(os.path.join(self.tmpdir, "nope")), []
        )


class TestDeleteDuplicatesSyncBody(AdminAPITestCase):

    def test_deletes_file_and_thumbnail_inside_output(self):
        target = self.output_dir / "sub" / "x.png"
        target.parent.mkdir()
        target.write_bytes(b"x")
        thumb = self.output_dir / "thumbnails" / "sub" / "x_thumb.png"
        thumb.parent.mkdir(parents=True)
        thumb.write_bytes(b"t")

        result = self.api._delete_duplicate_images_sync(
            [str(target)], str(self.output_dir)
        )

        self.assertEqual(result["deleted_count"], 1)
        self.assertEqual(result["failed_count"], 0)
        self.assertFalse(target.exists())
        self.assertFalse(thumb.exists())

    def test_refuses_paths_outside_output_and_reports_missing(self):
        secret = self.outside_dir / "secret.txt"
        secret.write_text("s")

        result = self.api._delete_duplicate_images_sync(
            [str(secret), "ghost.png"], str(self.output_dir)
        )

        self.assertEqual(result["deleted_count"], 0)
        self.assertEqual(result["failed_count"], 2)
        self.assertTrue(secret.exists())
        self.assertTrue(any("outside" in f for f in result["failed_files"]))
        self.assertTrue(any("not found" in f for f in result["failed_files"]))

    def test_without_output_dir_everything_fails(self):
        result = self.api._delete_duplicate_images_sync(["a.png"], None)
        self.assertEqual(result["failed_count"], 1)
        self.assertEqual(result["deleted_count"], 0)

    def test_refuses_non_media_files_inside_output(self):
        for name in ("notes.txt", "prompts.db", "script.py", "noext"):
            (self.output_dir / name).write_bytes(b"x")

        result = self.api._delete_duplicate_images_sync(
            ["notes.txt", "prompts.db", "script.py", "noext"], str(self.output_dir)
        )

        self.assertEqual(result["deleted_count"], 0)
        self.assertEqual(result["failed_count"], 4)
        for name in ("notes.txt", "prompts.db", "script.py", "noext"):
            self.assertTrue((self.output_dir / name).exists(), name)
        self.assertTrue(all("not a media file" in f for f in result["failed_files"]))

    def test_refuses_files_under_thumbnails(self):
        thumb = self.output_dir / "thumbnails" / "sub" / "x_thumb.png"
        thumb.parent.mkdir(parents=True)
        thumb.write_bytes(b"t")

        result = self.api._delete_duplicate_images_sync(
            [os.path.join("thumbnails", "sub", "x_thumb.png")], str(self.output_dir)
        )

        self.assertEqual(result["deleted_count"], 0)
        self.assertTrue(thumb.exists())
        self.assertIn("thumbnail", result["failed_files"][0])

    def test_every_scanned_media_extension_is_deletable(self):
        from py.api.admin import (
            DELETABLE_MEDIA_EXTENSIONS,
            IMAGE_EXTENSIONS,
            VIDEO_EXTENSIONS,
        )

        for ext in IMAGE_EXTENSIONS + VIDEO_EXTENSIONS + (".tif",):
            self.assertIn(ext, DELETABLE_MEDIA_EXTENSIONS)
        for ext in (".txt", ".db", ".py", ".json", ".html", ""):
            self.assertNotIn(ext, DELETABLE_MEDIA_EXTENSIONS)
        (self.output_dir / "a.WEBM").write_bytes(b"v")
        result = self.api._delete_duplicate_images_sync(
            ["a.WEBM"], str(self.output_dir)
        )
        self.assertEqual(result["deleted_count"], 1)


class TestCollectMediaFiles(AdminAPITestCase):
    """_collect_media_files: bounded walk that never follows symlinked files."""

    def _collect(self):
        from py.api.admin import (
            IMAGE_EXTENSIONS,
            VIDEO_EXTENSIONS,
            _collect_media_files,
        )

        found = _collect_media_files(
            [self.output_dir], IMAGE_EXTENSIONS + VIDEO_EXTENSIONS
        )
        return sorted(p.relative_to(self.output_dir).as_posix() for p in found)

    def test_matches_case_insensitively_and_skips_thumbnails(self):
        (self.output_dir / "a.png").write_bytes(b"x")
        (self.output_dir / "b.PNG").write_bytes(b"x")
        (self.output_dir / "c.Mp4").write_bytes(b"x")
        (self.output_dir / "d.txt").write_bytes(b"x")
        (self.output_dir / "sub").mkdir()
        (self.output_dir / "sub" / "e.jpg").write_bytes(b"x")
        (self.output_dir / "thumbnails").mkdir()
        (self.output_dir / "thumbnails" / "a_thumb.png").write_bytes(b"x")
        (self.output_dir / "sub" / "thumbnails").mkdir()
        (self.output_dir / "sub" / "thumbnails" / "e_thumb.jpg").write_bytes(b"x")
        self.assertEqual(self._collect(), ["a.png", "b.PNG", "c.Mp4", "sub/e.jpg"])

    def test_symlinked_files_and_directories_are_skipped(self):
        secret = Path(self.tmpdir) / "secret.png"
        secret.write_bytes(b"s")
        secret_dir = Path(self.tmpdir) / "secret_dir"
        secret_dir.mkdir()
        (secret_dir / "inner.png").write_bytes(b"s")
        (self.output_dir / "real.png").write_bytes(b"r")
        try:
            os.symlink(secret, self.output_dir / "link.png")
            os.symlink(secret_dir, self.output_dir / "linkdir")
        except (OSError, NotImplementedError):
            self.skipTest("symlinks not available")
        self.assertEqual(self._collect(), ["real.png"])

    def test_depth_is_capped(self):
        from py.api.admin import MAX_SCAN_DEPTH

        self.assertEqual(MAX_SCAN_DEPTH, 12)
        deep = self.output_dir
        for level in range(1, MAX_SCAN_DEPTH + 2):
            deep = deep / f"d{level}"
            deep.mkdir()
            (deep / f"f{level}.png").write_bytes(b"x")
        found = self._collect()
        self.assertIn(
            "/".join(f"d{i}" for i in range(1, MAX_SCAN_DEPTH + 1))
            + f"/f{MAX_SCAN_DEPTH}.png",
            found,
        )
        self.assertFalse(
            any(name.endswith(f"f{MAX_SCAN_DEPTH + 1}.png") for name in found)
        )
        self.assertEqual(len(found), MAX_SCAN_DEPTH)

    def test_file_count_is_capped(self):
        import py.api.admin as admin_module
        from unittest.mock import patch

        self.assertEqual(admin_module.MAX_SCAN_FILES, 50_000)
        for i in range(6):
            (self.output_dir / f"{i}.png").write_bytes(b"x")
        with patch.object(admin_module, "MAX_SCAN_FILES", 4):
            self.assertEqual(len(self._collect()), 4)

    def test_missing_root_is_skipped(self):
        from py.api.admin import _collect_media_files

        self.assertEqual(
            _collect_media_files([Path(self.tmpdir) / "nope"], (".png",)), []
        )


class TestSettingsReadMonitorThroughModule(AdminAPITestCase):
    """get_settings reaches the image monitor through utils.image_monitor."""

    def test_no_sys_modules_scan(self):
        import inspect

        from py.api.admin import AdminRoutesMixin

        source = inspect.getsource(AdminRoutesMixin.get_settings)
        self.assertNotIn("sys.modules", source)

    async def test_reports_directories_of_running_monitor(self):
        import utils.image_monitor as im_mod
        from utils.image_monitor import get_image_monitor

        orig = im_mod._monitor_instance
        im_mod._monitor_instance = None
        self.addCleanup(setattr, im_mod, "_monitor_instance", orig)
        monitor = get_image_monitor(MagicMock(), MagicMock())
        monitor.monitored_directories = [self._root(self.output_dir)]

        resp = await self.client.request("GET", "/prompt_manager/settings")

        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertEqual(data["settings"]["monitored_directories"], ["output"])

    async def test_falls_back_to_configured_roots_without_monitor(self):
        import utils.image_monitor as im_mod

        orig = im_mod._monitor_instance
        im_mod._monitor_instance = None
        self.addCleanup(setattr, im_mod, "_monitor_instance", orig)
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

        resp = await self.client.request("GET", "/prompt_manager/settings")

        data = await resp.json()
        self.assertEqual(data["settings"]["monitored_directories"], ["output"])
        self.assertEqual(data["settings"]["gallery_root_paths"], ["output"])


def _raise(*args, **kwargs):
    raise RuntimeError("boom")


class TestStatsCleanupAndErrors(AdminAPITestCase):

    async def test_stats_success(self):
        resp = await self.client.request("GET", "/prompt_manager/stats")
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertIn("total_prompts", data["stats"])

    async def test_stats_db_failure_is_500(self):
        self.api.db.get_statistics = _raise
        resp = await self.client.request("GET", "/prompt_manager/stats")
        self.assertEqual(resp.status, 500)
        self.assertFalse((await resp.json())["success"])

    async def test_cleanup_success(self):
        resp = await self.client.request("POST", "/prompt_manager/cleanup")
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertEqual(data["duplicates_removed"], 0)

    async def test_cleanup_failure_is_500(self):
        self.api.db.cleanup_duplicates = _raise
        resp = await self.client.request("POST", "/prompt_manager/cleanup")
        self.assertEqual(resp.status, 500)

    async def test_scan_duplicates_failure_is_500(self):
        self.api._find_comfyui_output_dir = _raise
        resp = await self.client.request("GET", "/prompt_manager/scan_duplicates")
        body = await resp.text()
        self.assertEqual(resp.status, 500)
        self.assertNotIn("boom", body)

    async def test_scan_duplicates_without_output_dir_is_empty(self):
        self.api._find_comfyui_output_dir = lambda: None
        resp = await self.client.request("GET", "/prompt_manager/scan_duplicates")
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertEqual(data["duplicates"], [])


class TestDeleteDuplicatesEndpoint(AdminAPITestCase):

    async def _post(self, payload):
        return await self.client.request(
            "POST", "/prompt_manager/delete_duplicate_images", json=payload
        )

    async def test_empty_list_is_400(self):
        resp = await self._post({"image_paths": []})
        self.assertEqual(resp.status, 400)

    async def test_non_list_is_400(self):
        resp = await self._post({"image_paths": "a.png"})
        self.assertEqual(resp.status, 400)

    async def test_malformed_json_is_400(self):
        resp = await self.client.request(
            "POST",
            "/prompt_manager/delete_duplicate_images",
            data=b"{",
            headers={"Content-Type": "application/json"},
        )
        self.assertEqual(resp.status, 400)

    async def test_failures_are_reported(self):
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        resp = await self._post({"image_paths": ["ghost.png"]})
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertEqual(data["failed_count"], 1)
        self.assertIn("1 failed", data["message"])
        self.assertEqual(len(data["failed_files"]), 1)

    async def test_unexpected_error_is_500(self):
        self.api._find_comfyui_output_dir = _raise
        resp = await self._post({"image_paths": ["a.png"]})
        self.assertEqual(resp.status, 500)

    def test_thumbnail_removal_failure_does_not_fail_delete(self):
        target = self.output_dir / "x.png"
        target.write_bytes(b"x")
        thumb_dir = self.output_dir / "thumbnails" / "x_thumb.png"
        thumb_dir.mkdir(parents=True)
        (thumb_dir / "child").write_bytes(b"c")  # os.remove on a dir raises OSError

        result = self.api._delete_duplicate_images_sync(
            [str(target)], str(self.output_dir)
        )

        self.assertEqual(result["deleted_count"], 1)
        self.assertFalse(target.exists())
        self.assertTrue(thumb_dir.exists())

    def test_unexpected_error_per_file_is_reported_without_path(self):
        self.api._delete_one_duplicate = lambda *a: (_ for _ in ()).throw(
            PermissionError(13, "Permission denied", os.path.join(self.tmpdir, "x.png"))
        )
        result = self.api._delete_duplicate_images_sync(["x.png"], str(self.output_dir))
        self.assertEqual(result["failed_count"], 1)
        self.assertIn("Permission denied", result["failed_files"][0])
        self.assertNotIn(self.tmpdir, result["failed_files"][0])


class TestJsonBodyCap(AdminAPITestCase):
    """Every JSON POST in the admin API rejects bodies over the shared cap."""

    JSON_POSTS = (
        "/prompt_manager/settings",
        "/prompt_manager/delete_duplicate_images",
        "/prompt_manager/diagnostics/test-link",
        "/prompt_manager/maintenance",
    )

    async def test_oversized_bodies_are_413(self):
        body = json.dumps({"pad": "x" * 1_000_100}).encode()
        for path in self.JSON_POSTS:
            resp = await self.client.request(
                "POST", path, data=body, headers={"Content-Type": "application/json"}
            )
            self.assertEqual(resp.status, 413, path)
            self.assertFalse((await resp.json())["success"], path)


class TestDuplicateJobsAreSingleFlight(AdminAPITestCase):
    """A duplicate scan or delete refuses to overlap with a running one."""

    def _block_scan(self):
        gate = asyncio.Event()

        async def slow_scan():
            await gate.wait()
            return []

        self.api.find_duplicate_images = slow_scan
        return gate

    async def test_second_scan_while_one_runs_is_409(self):
        gate = self._block_scan()
        first = asyncio.ensure_future(
            self.client.request("GET", "/prompt_manager/scan_duplicates")
        )
        await asyncio.sleep(0.05)

        second = await asyncio.wait_for(
            self.client.request("GET", "/prompt_manager/scan_duplicates"), 5
        )
        self.assertEqual(second.status, 409)
        data = await second.json()
        self.assertFalse(data["success"])
        self.assertIn("already running", data["error"])

        gate.set()
        self.assertEqual((await first).status, 200)

    async def test_delete_while_scan_runs_is_409_and_scan_after_is_fine(self):
        gate = self._block_scan()
        first = asyncio.ensure_future(
            self.client.request("GET", "/prompt_manager/scan_duplicates")
        )
        await asyncio.sleep(0.05)

        resp = await asyncio.wait_for(
            self.client.request(
                "POST",
                "/prompt_manager/delete_duplicate_images",
                json={"image_paths": ["a.png"]},
            ),
            5,
        )
        self.assertEqual(resp.status, 409)

        gate.set()
        await first
        resp = await self.client.request("GET", "/prompt_manager/scan_duplicates")
        self.assertEqual(resp.status, 200)

    async def test_lock_is_released_after_a_failed_scan(self):
        self.api._find_comfyui_output_dir = _raise
        resp = await self.client.request("GET", "/prompt_manager/scan_duplicates")
        self.assertEqual(resp.status, 500)
        self.api._find_comfyui_output_dir = lambda: None
        resp = await self.client.request("GET", "/prompt_manager/scan_duplicates")
        self.assertEqual(resp.status, 200)


class TestSettingsMisc(AdminAPITestCase):

    async def test_invalid_payload_changes_nothing(self):
        before = (
            PromptManagerConfig.RESULT_TIMEOUT,
            PromptManagerConfig.WEBUI_DISPLAY_MODE,
            PromptManagerConfig.WORKER_THREADS,
            PromptManagerConfig.INFINITE_SCROLL,
            list(GalleryConfig.MONITORING_DIRECTORIES),
        )

        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={
                "result_timeout": 42,
                "webui_display_mode": "popup",
                "worker_threads": 1,
                "infinite_scroll": True,
                "gallery_root_paths": [os.path.abspath(os.sep)],  # rejected
            },
        )

        self.assertEqual(resp.status, 400)
        after = (
            PromptManagerConfig.RESULT_TIMEOUT,
            PromptManagerConfig.WEBUI_DISPLAY_MODE,
            PromptManagerConfig.WORKER_THREADS,
            PromptManagerConfig.INFINITE_SCROLL,
            list(GalleryConfig.MONITORING_DIRECTORIES),
        )
        self.assertEqual(after, before)
        self.assertFalse(os.path.exists(self.config_path))

    async def test_infinite_scroll_default_is_reported_saved_and_validated(self):
        data = await (
            await self.client.request("GET", "/prompt_manager/settings")
        ).json()
        self.assertIs(data["settings"]["infinite_scroll"], False)

        resp = await self.client.request(
            "POST", "/prompt_manager/settings", json={"infinite_scroll": True}
        )
        self.assertEqual(resp.status, 200)
        self.assertTrue(PromptManagerConfig.INFINITE_SCROLL)
        self.assertIs(self._read_config()["web_ui"]["infinite_scroll"], True)
        data = await (
            await self.client.request("GET", "/prompt_manager/settings")
        ).json()
        self.assertIs(data["settings"]["infinite_scroll"], True)

        resp = await self.client.request(
            "POST", "/prompt_manager/settings", json={"infinite_scroll": "on"}
        )
        self.assertEqual(resp.status, 400)
        self.assertTrue(PromptManagerConfig.INFINITE_SCROLL)

    async def test_result_timeout_and_display_mode_saved(self):
        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"result_timeout": 9, "webui_display_mode": "popup"},
        )
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertFalse(data["restart_required"])
        self.assertEqual(PromptManagerConfig.RESULT_TIMEOUT, 9)
        self.assertEqual(PromptManagerConfig.WEBUI_DISPLAY_MODE, "popup")
        saved = self._read_config()
        self.assertEqual(saved["web_ui"]["result_timeout"], 9)

    async def test_save_merges_into_existing_config_file(self):
        self._write_config(
            {
                "integrations": {"lora_manager": {"enabled": True, "path": "x"}},
                "database": {"default_path": "custom.db"},
                "web_ui": {"result_timeout": 1, "show_test_button": True},
            }
        )

        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={
                "result_timeout": 7,
                "gallery_root_paths": [self._root(self.output_dir)],
            },
        )

        self.assertEqual(resp.status, 200)
        saved = self._read_config()
        self.assertEqual(saved["integrations"]["lora_manager"]["enabled"], True)
        self.assertEqual(saved["integrations"]["lora_manager"]["path"], "x")
        self.assertEqual(saved["database"]["default_path"], "custom.db")
        self.assertEqual(saved["web_ui"]["result_timeout"], 7)
        self.assertEqual(saved["web_ui"]["show_test_button"], True)
        self.assertEqual(
            saved["gallery"]["monitoring"]["directories"], [self._root(self.output_dir)]
        )

    async def test_save_replaces_a_corrupt_config_file(self):
        with open(self.config_path, "w") as f:
            f.write("{not json")
        resp = await self.client.request(
            "POST", "/prompt_manager/settings", json={"result_timeout": 4}
        )
        self.assertEqual(resp.status, 200)
        self.assertEqual(self._read_config()["web_ui"]["result_timeout"], 4)

    async def test_saved_config_file_is_private_to_the_user(self):
        import stat

        resp = await self.client.request(
            "POST", "/prompt_manager/settings", json={"result_timeout": 4}
        )
        self.assertEqual(resp.status, 200)
        if os.name == "posix":
            mode = stat.S_IMODE(os.stat(self.config_path).st_mode)
            self.assertEqual(mode, 0o600)

    async def test_invalid_result_timeout_is_400(self):
        resp = await self.client.request(
            "POST", "/prompt_manager/settings", json={"result_timeout": "soon"}
        )
        self.assertEqual(resp.status, 400)

    async def test_malformed_json_is_400(self):
        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            data=b"{oops",
            headers={"Content-Type": "application/json"},
        )
        self.assertEqual(resp.status, 400)

    async def test_unchanged_roots_do_not_require_restart(self):
        GalleryConfig.MONITORING_DIRECTORIES = [
            os.path.normcase(os.path.realpath(self.output_dir))
        ]
        resp = await self.client.request(
            "POST",
            "/prompt_manager/settings",
            json={"gallery_root_paths": [self._root(self.output_dir)]},
        )
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertFalse(data["restart_required"])

    async def test_unwritable_config_path_still_succeeds(self):
        os.environ[CONFIG_PATH_ENV] = os.path.join(self.tmpdir, "missing", "c.json")
        resp = await self.client.request(
            "POST", "/prompt_manager/settings", json={"result_timeout": 3}
        )
        self.assertEqual(resp.status, 200)

    async def test_save_unexpected_error_is_500(self):
        self.api._parse_gallery_roots = _raise
        resp = await self.client.request(
            "POST", "/prompt_manager/settings", json={"gallery_root_paths": []}
        )
        self.assertEqual(resp.status, 500)
        self.assertNotIn("boom", await resp.text())

    async def test_get_unexpected_error_is_500(self):
        self.api._monitored_directories = _raise
        resp = await self.client.request("GET", "/prompt_manager/settings")
        self.assertEqual(resp.status, 500)


class TestDiagnostics(AdminAPITestCase):

    def _install_monitor(self, monitor):
        import utils.image_monitor as im_mod

        orig = im_mod._monitor_instance
        im_mod._monitor_instance = monitor
        self.addCleanup(setattr, im_mod, "_monitor_instance", orig)

    async def _diag(self):
        resp = await self.client.request("GET", "/prompt_manager/diagnostics")
        return resp, (await resp.json())["diagnostics"]

    async def test_healthy_report(self):
        self.api.db.save_prompt(text="hello", prompt_hash="h1")
        self._install_monitor(FakeMonitor([self._root(self.output_dir)]))

        resp, diag = await self._diag()

        self.assertEqual(resp.status, 200)
        self.assertEqual(diag["database"]["status"], "ok")
        self.assertEqual(diag["database"]["prompt_count"], 1)
        self.assertTrue(diag["database"]["has_images_table"])
        self.assertEqual(diag["dependencies"]["status"], "ok")
        self.assertEqual(diag["comfyui_output"]["status"], "ok")
        self.assertEqual(diag["image_monitor"]["status"], "ok")

    async def test_corrupt_database_reported(self):
        bad = os.path.join(self.tmpdir, "bad.db")
        with open(bad, "wb") as f:
            f.write(b"not a database at all, definitely not sqlite")
        self.api.db.model.db_path = bad

        _, diag = await self._diag()

        self.assertEqual(diag["database"]["status"], "error")
        self.assertIn("Database error", diag["database"]["message"])

    async def test_monitor_not_running_and_no_output_dirs(self):
        self._install_monitor(None)
        sys.modules.pop("folder_paths", None)
        orig_cwd = os.getcwd()
        os.chdir(self.outside_dir)
        self.addCleanup(os.chdir, orig_cwd)

        _, diag = await self._diag()

        self.assertEqual(diag["image_monitor"]["status"], "error")
        self.assertIn("not initialized", diag["image_monitor"]["message"])
        self.assertEqual(diag["comfyui_output"]["status"], "warning")
        self.assertEqual(diag["comfyui_output"]["output_dirs"], [])

    async def test_relative_output_fallback_without_folder_paths(self):
        sys.modules.pop("folder_paths", None)
        orig_cwd = os.getcwd()
        os.chdir(self.comfy_dir)  # has an "output" child
        self.addCleanup(os.chdir, orig_cwd)

        _, diag = await self._diag()

        self.assertEqual(diag["comfyui_output"]["status"], "ok")
        self.assertEqual(diag["comfyui_output"]["output_dirs"], ["output"])

    async def test_monitor_status_failure_reported(self):
        broken = MagicMock()
        broken.get_status = _raise
        self._install_monitor(broken)

        _, diag = await self._diag()

        self.assertEqual(diag["image_monitor"]["status"], "error")
        self.assertIn("Failed to get monitor status", diag["image_monitor"]["message"])

    async def test_unexpected_error_is_500(self):
        import py.api.admin as admin_module

        orig = admin_module._diagnose_database
        admin_module._diagnose_database = _raise
        self.addCleanup(setattr, admin_module, "_diagnose_database", orig)
        resp = await self.client.request("GET", "/prompt_manager/diagnostics")
        self.assertEqual(resp.status, 500)


class TestTestImageLink(AdminAPITestCase):

    async def _post(self, payload):
        return await self.client.request(
            "POST", "/prompt_manager/diagnostics/test-link", json=payload
        )

    async def test_missing_prompt_id_is_400(self):
        resp = await self._post({})
        self.assertEqual(resp.status, 400)

    async def test_links_test_image(self):
        pid = self.api.db.save_prompt(text="hello", prompt_hash="h1")
        resp = await self._post({"prompt_id": pid})
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertEqual(data["result"]["status"], "ok")
        self.assertIsInstance(data["result"]["image_id"], int)

    async def test_client_supplied_path_is_not_stored(self):
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        pid = self.api.db.save_prompt(text="hello", prompt_hash="h1")
        evil = os.path.join(self.tmpdir, "etc", "passwd")

        resp = await self._post({"prompt_id": pid, "image_path": evil})

        data = await resp.json()
        self.assertEqual(resp.status, 200, data)
        images = self.api.db.get_prompt_images(pid)
        self.assertEqual(len(images), 1)
        stored = Path(images[0]["image_path"])
        self.assertNotEqual(os.path.normcase(str(stored)), os.path.normcase(evil))
        self.assertTrue(
            stored.resolve().is_relative_to(self.output_dir.resolve()),
            stored,
        )
        self.assertIn("test", stored.name.lower())
        self.assertNotIn(self.tmpdir, json.dumps(data))

    async def test_without_output_dir_a_relative_marker_is_used(self):
        self.api._find_comfyui_output_dir = lambda: None
        pid = self.api.db.save_prompt(text="hello", prompt_hash="h1")
        resp = await self._post({"prompt_id": pid, "image_path": "/x/y.png"})
        self.assertEqual(resp.status, 200)
        images = self.api.db.get_prompt_images(pid)
        self.assertFalse(os.path.isabs(images[0]["image_path"]))

    async def test_db_failure_is_reported_in_result(self):
        self.api.db.link_image_to_prompt = _raise
        resp = await self._post({"prompt_id": 1})
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertFalse(data["success"])
        self.assertEqual(data["result"]["status"], "error")

    async def test_malformed_json_is_400(self):
        resp = await self.client.request(
            "POST",
            "/prompt_manager/diagnostics/test-link",
            data=b"nope",
            headers={"Content-Type": "application/json"},
        )
        self.assertEqual(resp.status, 400)


class TestMaintenance(AdminAPITestCase):

    ALL_OPS = [
        "cleanup_duplicates",
        "vacuum",
        "cleanup_orphaned_images",
        "check_hash_duplicates",
        "statistics",
        "prune_orphaned_prompts",
        "check_consistency",
    ]

    async def _post(self, payload=None, **kwargs):
        if payload is None:
            return await self.client.request(
                "POST", "/prompt_manager/maintenance", **kwargs
            )
        return await self.client.request(
            "POST", "/prompt_manager/maintenance", json=payload, **kwargs
        )

    async def test_default_operations_without_body(self):
        resp = await self._post(data=b"", headers={"Content-Type": "text/plain"})
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertEqual(data["operations_completed"], 3)
        self.assertTrue(data["all_successful"])
        self.assertEqual(
            set(data["results"]),
            {"cleanup_duplicates", "vacuum", "cleanup_orphaned_images"},
        )

    async def test_all_operations(self):
        self.api.db.save_prompt(text="orphan", prompt_hash="h1")
        resp = await self._post({"operations": self.ALL_OPS})
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertEqual(data["operations_completed"], len(self.ALL_OPS))
        self.assertTrue(data["all_successful"], data["results"])
        self.assertIn("info", data["results"]["statistics"])
        self.assertEqual(data["results"]["check_consistency"]["issues_found"], 0)

    async def test_statistics_database_path_is_public(self):
        resp = await self._post({"operations": ["statistics"]})
        body = await resp.text()
        self.assertEqual(resp.status, 200)
        self.assertNotIn(self.tmpdir, body)
        info = json.loads(body)["results"]["statistics"]["info"]
        self.assertEqual(info["database_path"], "prompts.db")

    async def test_unknown_operation_does_nothing(self):
        resp = await self._post({"operations": ["reboot"]})
        data = await resp.json()
        self.assertEqual(data["operations_completed"], 0)
        self.assertTrue(data["all_successful"])

    async def test_each_operation_reports_its_own_failure(self):
        self.api.db.cleanup_duplicates = _raise
        self.api.db.model.vacuum_database = _raise
        self.api.db.cleanup_missing_images = _raise
        self.api.db.check_hash_duplicates = _raise
        self.api.db.model.get_database_info = _raise
        self.api.db.prune_orphaned_prompts = _raise
        self.api.db.check_consistency = _raise

        resp = await self._post({"operations": self.ALL_OPS})
        data = await resp.json()

        self.assertEqual(resp.status, 200)
        self.assertFalse(data["all_successful"])
        self.assertEqual(len(data["results"]), len(self.ALL_OPS))
        for name in self.ALL_OPS:
            self.assertFalse(data["results"][name]["success"], name)
            self.assertEqual(data["results"][name]["error"], "boom")

    async def test_unexpected_error_is_500(self):
        self.api._run_in_executor = _raise
        resp = await self._post({"operations": ["vacuum"]})
        self.assertEqual(resp.status, 500)


class TestBackupRestoreContract(AdminAPITestCase):
    """Contract tests only: status codes and envelope of backup/restore."""

    async def test_backup_downloads_sqlite_file(self):
        self.api.db.save_prompt(text="keep me", prompt_hash="h1")
        resp = await self.client.request("GET", "/prompt_manager/backup")
        body = await resp.read()
        self.assertEqual(resp.status, 200)
        self.assertIn("attachment", resp.headers["Content-Disposition"])
        self.assertTrue(body.startswith(b"SQLite format 3"))

    async def test_backup_without_db_file_is_404(self):
        self.api.db.model.db_path = os.path.join(self.tmpdir, "nope.db")
        resp = await self.client.request("GET", "/prompt_manager/backup")
        self.assertEqual(resp.status, 404)

    async def _restore(self, field_name, content):
        from aiohttp import FormData

        form = FormData()
        form.add_field(field_name, content, filename="upload.db")
        return await self.client.request("POST", "/prompt_manager/restore", data=form)

    async def test_restore_wrong_field_is_400(self):
        resp = await self._restore("other", b"x")
        self.assertEqual(resp.status, 400)

    async def test_restore_empty_file_is_400(self):
        resp = await self._restore("database_file", b"")
        self.assertEqual(resp.status, 400)

    async def test_restore_non_sqlite_is_400_and_db_untouched(self):
        self.api.db.save_prompt(text="keep me", prompt_hash="h1")
        resp = await self._restore("database_file", b"definitely not sqlite data")
        self.assertEqual(resp.status, 400)
        self.assertIsNotNone(self.api.db.get_prompt_by_hash("h1"))

    def _snapshot_upload(self, extra_sql=()):
        """Bytes of a self-contained copy of the live DB (WAL rows included)."""
        import sqlite3
        from contextlib import closing

        snapshot = os.path.join(self.tmpdir, "snapshot.db")
        with closing(sqlite3.connect(self.api.db.model.db_path)) as src:
            with closing(sqlite3.connect(snapshot)) as dst:
                src.backup(dst)
                for statement in extra_sql:
                    dst.execute(statement)
                dst.commit()
        with open(snapshot, "rb") as f:
            return f.read()

    async def test_restore_valid_db_succeeds(self):
        import sqlite3
        from contextlib import closing

        self.api.db.save_prompt(text="restored", prompt_hash="h1")
        upload = self._snapshot_upload()
        resp = await self._restore("database_file", upload)
        data = await resp.json()
        self.assertEqual(resp.status, 200, data)
        self.assertTrue(data["success"])
        self.assertEqual(data["prompt_count"], 1)
        with closing(sqlite3.connect(self.api.db.model.db_path)) as conn:
            self.assertEqual(
                conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0], 1
            )

    async def test_restore_backup_created_is_a_public_path(self):
        self.api.db.save_prompt(text="restored", prompt_hash="h1")
        resp = await self._restore("database_file", self._snapshot_upload())
        body = await resp.text()
        self.assertEqual(resp.status, 200, body)
        self.assertNotIn(self.tmpdir, body)
        data = json.loads(body)
        self.assertTrue(data["backup_created"].startswith("prompts.db.backup_"))

    async def test_restore_rejects_database_with_a_trigger(self):
        self.api.db.save_prompt(text="keep me", prompt_hash="h1")
        upload = self._snapshot_upload(
            [
                "CREATE TRIGGER evil AFTER INSERT ON prompts "
                "BEGIN DELETE FROM prompts; END"
            ]
        )
        resp = await self._restore("database_file", upload)
        data = await resp.json()
        self.assertEqual(resp.status, 400, data)
        self.assertIn("trigger", data["error"].lower())
        self.assertIsNotNone(self.api.db.get_prompt_by_hash("h1"))

    async def test_restore_rejects_database_with_a_view(self):
        self.api.db.save_prompt(text="keep me", prompt_hash="h1")
        upload = self._snapshot_upload(["CREATE VIEW peek AS SELECT text FROM prompts"])
        resp = await self._restore("database_file", upload)
        data = await resp.json()
        self.assertEqual(resp.status, 400, data)
        self.assertIn("view", data["error"].lower())

    async def test_restore_verifies_in_the_executor(self):
        self.api.db.save_prompt(text="restored", prompt_hash="h1")
        upload = self._snapshot_upload()
        calls = []
        original = self.api._run_in_executor

        async def spy(func, *args, **kwargs):
            calls.append(getattr(func, "__name__", repr(func)))
            return await original(func, *args, **kwargs)

        self.api._run_in_executor = spy
        resp = await self._restore("database_file", upload)
        self.assertEqual(resp.status, 200)
        self.assertIn("verify_database_file", calls)
        self.assertIn("restore_from_file", calls)

    async def test_restore_keeps_at_most_ten_safety_backups(self):
        import time

        from py.api.admin import MAX_SAFETY_BACKUPS

        self.assertEqual(MAX_SAFETY_BACKUPS, 10)
        db_path = self.api.db.model.db_path
        stale = []
        now = time.time()
        for i in range(12):
            path = f"{db_path}.backup_202001{i + 1:02d}_000000"
            with open(path, "wb") as f:
                f.write(b"old")
            os.utime(path, (now - 10_000 + i, now - 10_000 + i))
            stale.append(path)
        unrelated = os.path.join(self.tmpdir, "other.db.backup_20200101_000000")
        with open(unrelated, "wb") as f:
            f.write(b"x")
        self.api.db.save_prompt(text="restored", prompt_hash="h1")

        resp = await self._restore("database_file", self._snapshot_upload())

        self.assertEqual(resp.status, 200)
        remaining = sorted(
            name
            for name in os.listdir(self.tmpdir)
            if name.startswith("prompts.db.backup_")
        )
        self.assertEqual(len(remaining), MAX_SAFETY_BACKUPS)
        data = await resp.json()
        self.assertIn(data["backup_created"], remaining)
        for path in stale[:3]:
            self.assertFalse(os.path.exists(path), path)
        self.assertTrue(os.path.exists(unrelated))

    async def test_restore_failure_body_has_no_absolute_path(self):
        self.api.db.save_prompt(text="restored", prompt_hash="h1")
        upload = self._snapshot_upload()

        def boom(src_path):
            raise PermissionError(
                13, "Permission denied", os.path.join(self.tmpdir, "prompts.db")
            )

        self.api.db.model.restore_from_file = boom
        resp = await self._restore("database_file", upload)
        body = await resp.text()
        self.assertEqual(resp.status, 500)
        self.assertNotIn(self.tmpdir, body)
        self.assertIn("Permission denied", json.loads(body)["error"])

    async def test_backup_failure_body_has_no_absolute_path(self):
        def boom(path):
            raise PermissionError(
                13, "Permission denied", os.path.join(self.tmpdir, "x.db")
            )

        self.api.db.model.backup_database = boom
        resp = await self.client.request("GET", "/prompt_manager/backup")
        body = await resp.text()
        self.assertEqual(resp.status, 500)
        self.assertNotIn(self.tmpdir, body)
        self.assertIn("Permission denied", json.loads(body)["error"])


class TestOutputScanInternals(AdminAPITestCase):

    def test_metadata_extraction_failure_yields_empty(self):
        self.api._extract_comfyui_metadata = _raise
        results = self.api._extract_batch_metadata_sync([Path("a.png")])
        self.assertEqual(results, [(Path("a.png"), {})])

    def test_ingest_outcomes(self):
        f = Path("gen.png")
        meta = {
            "prompt": json.dumps(
                {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}}
            )
        }
        self.assertIsNone(self.api._ingest_scanned_file(f, {}))
        self.assertIsNone(self.api._ingest_scanned_file(f, {"prompt": "{}"}))
        self.assertEqual(
            self.api._ingest_scanned_file(f, {"parameters": "  \nNegative prompt: x"}),
            "found",
        )
        self.assertEqual(self.api._ingest_scanned_file(f, meta), "added")
        self.assertEqual(self.api._ingest_scanned_file(f, meta), "linked")

    def test_ingest_link_failures_are_tolerated(self):
        f = Path("gen.png")
        meta = {
            "prompt": json.dumps(
                {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a dog"}}}
            )
        }
        self.api.db.link_image_to_prompt = _raise
        self.assertEqual(self.api._ingest_scanned_file(f, meta), "added")
        self.assertEqual(self.api._ingest_scanned_file(f, meta), "found")

    def test_ingest_save_failure_counts_as_found(self):
        f = Path("gen.png")
        meta = {
            "prompt": json.dumps(
                {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a bird"}}}
            )
        }
        self.api.db.save_prompt = lambda *a, **k: None
        self.assertEqual(self.api._ingest_scanned_file(f, meta), "found")

    def test_batch_counts_and_per_file_errors(self):
        meta = {
            "prompt": json.dumps(
                {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a fish"}}}
            )
        }
        self.api.db.get_prompt_by_hash = _raise
        counts = self.api._ingest_scan_batch_sync(
            [(Path("a.png"), meta), (Path("b.png"), {})]
        )
        self.assertEqual(counts, {"processed": 2, "found": 0, "added": 0, "linked": 0})

    async def test_scan_internal_error_streams_error_event(self):
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        self.api._collect_output_media_sync = _raise
        body = await (await self.client.request("POST", "/prompt_manager/scan")).text()
        self.assertIn('"type": "error"', body)
        self.assertNotIn("boom", body)


class TestOutputScanIsABackgroundJob(AdminAPITestCase):
    """The output scan keeps running when the browser goes away, and a second
    request attaches to the running scan instead of starting another one."""

    def _gate_collect(self):
        """Block the media collection step until the returned event is set."""
        import threading

        gate = threading.Event()
        calls = []
        original = self.api._collect_output_media_sync

        def gated(output_dirs):
            calls.append(output_dirs)
            gate.wait(5)
            return original(output_dirs)

        self.api._collect_output_media_sync = gated
        return gate, calls

    async def test_scan_survives_client_disconnect_and_finishes(self):
        from unittest.mock import AsyncMock, patch

        _write_png(self.output_dir / "gen.png", prompt_text="a red square")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

        with self.assertNoLogs("aiohttp.server", level="ERROR"):
            with patch.object(
                web.StreamResponse,
                "write",
                AsyncMock(side_effect=ConnectionResetError("closing transport")),
            ):
                resp = await self.client.request("POST", "/prompt_manager/scan")
                self.assertEqual(resp.status, 200)
                await resp.release()
            await asyncio.wait_for(self.api._scan_job.task, 10)

        prompts = self.api.db.get_recent_prompts(limit=10)["prompts"]
        self.assertEqual([p["text"] for p in prompts], ["a red square"])
        self.assertEqual(self.api._scan_job.last_event["type"], "complete")
        self.assertFalse(self.api._scan_job.running)

    async def test_second_request_attaches_to_the_running_scan(self):
        _write_png(self.output_dir / "gen.png", prompt_text="a red square")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        gate, calls = self._gate_collect()

        first = asyncio.ensure_future(
            self.client.request("POST", "/prompt_manager/scan")
        )
        await asyncio.sleep(0.05)
        second = await asyncio.wait_for(
            self.client.request("POST", "/prompt_manager/scan"), 5
        )
        self.assertEqual(second.status, 200)
        self.assertIn("text/event-stream", second.headers["Content-Type"])

        gate.set()
        first_body = await (await first).text()
        second_body = await second.text()

        self.assertEqual(len(calls), 1)
        self.assertIn('"type": "complete"', first_body)
        self.assertIn('"type": "complete"', second_body)
        self.assertIn('"added": 1', second_body)

    async def test_status_reports_running_scan_and_its_last_event(self):
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        gate, _ = self._gate_collect()

        first = asyncio.ensure_future(
            self.client.request("POST", "/prompt_manager/scan")
        )
        await asyncio.sleep(0.05)
        resp = await self.client.request("GET", "/prompt_manager/scan/status")
        running = await resp.json()

        gate.set()
        await (await first).text()
        resp = await self.client.request("GET", "/prompt_manager/scan/status")
        finished = await resp.json()

        self.assertEqual(resp.status, 200)
        self.assertTrue(running["success"])
        self.assertTrue(running["running"])
        self.assertEqual(running["last_event"]["type"], "progress")
        self.assertFalse(finished["running"])
        self.assertEqual(finished["last_event"]["type"], "complete")

    async def test_status_before_any_scan_is_idle(self):
        resp = await self.client.request("GET", "/prompt_manager/scan/status")
        data = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertEqual(data, {"success": True, "running": False, "last_event": None})

    async def test_scan_after_a_finished_one_starts_fresh(self):
        _write_png(self.output_dir / "gen.png", prompt_text="a red square")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        gate, calls = self._gate_collect()
        gate.set()

        await (await self.client.request("POST", "/prompt_manager/scan")).text()
        body = await (await self.client.request("POST", "/prompt_manager/scan")).text()

        self.assertEqual(len(calls), 2)
        self.assertIn('"linked": 1', body)

    async def test_failed_scan_is_not_left_running(self):
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        self.api._collect_output_media_sync = _raise

        await (await self.client.request("POST", "/prompt_manager/scan")).text()
        data = await (
            await self.client.request("GET", "/prompt_manager/scan/status")
        ).json()

        self.assertFalse(data["running"])
        self.assertEqual(data["last_event"]["type"], "error")


class TestWorkerThreadsSetting(AdminAPITestCase):
    """The worker-thread count is exposed, validated, persisted and used by the scan."""

    async def test_get_settings_reports_worker_threads_and_cpu_count(self):
        data = await (
            await self.client.request("GET", "/prompt_manager/settings")
        ).json()
        settings = data["settings"]
        self.assertEqual(
            settings["cpu_count"], PromptManagerConfig.max_worker_threads()
        )
        self.assertEqual(settings["worker_threads"], PromptManagerConfig.WORKER_THREADS)
        self.assertGreaterEqual(settings["worker_threads"], 1)

    async def test_save_persists_worker_threads_under_performance(self):
        self._write_config({"performance": {"max_search_results": 25}})

        resp = await self.client.request(
            "POST", "/prompt_manager/settings", json={"worker_threads": 1}
        )

        self.assertEqual(resp.status, 200)
        self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
        saved = self._read_config()
        self.assertEqual(saved["performance"]["worker_threads"], 1)
        self.assertEqual(saved["performance"]["max_search_results"], 25)

    async def test_save_rejects_more_threads_than_cores(self):
        before = PromptManagerConfig.WORKER_THREADS
        too_many = PromptManagerConfig.max_worker_threads() + 1

        resp = await self.client.request(
            "POST", "/prompt_manager/settings", json={"worker_threads": too_many}
        )

        self.assertEqual(resp.status, 400)
        self.assertFalse((await resp.json())["success"])
        self.assertEqual(PromptManagerConfig.WORKER_THREADS, before)

    async def test_scan_reads_metadata_on_several_threads(self):
        import threading
        import time

        for i in range(8):
            _write_png(self.output_dir / f"gen{i}.png", prompt_text=f"prompt {i}")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        PromptManagerConfig.WORKER_THREADS = 4
        seen = set()
        lock = threading.Lock()
        original = self.api._extract_comfyui_metadata

        def recording(path):
            with lock:
                seen.add(threading.current_thread().name)
            time.sleep(0.05)
            return original(path)

        self.api._extract_comfyui_metadata = recording

        body = await (await self.client.request("POST", "/prompt_manager/scan")).text()

        self.assertIn('"added": 8', body)
        self.assertGreaterEqual(len(seen), 2)

    async def test_scan_with_one_worker_still_completes(self):
        _write_png(self.output_dir / "gen.png", prompt_text="solo")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        PromptManagerConfig.WORKER_THREADS = 1

        body = await (await self.client.request("POST", "/prompt_manager/scan")).text()

        self.assertIn('"added": 1', body)


def _write_mp4(path, comment):
    """Tiny black video with ``comment`` in its container tags (needs ffmpeg)."""
    import subprocess

    subprocess.run(
        [
            "ffmpeg",
            "-v",
            "quiet",
            "-f",
            "lavfi",
            "-i",
            "color=c=black:s=16x16:d=0.2",
            "-metadata",
            f"comment={comment}",
            "-y",
            str(path),
        ],
        check=True,
        timeout=60,
    )


class TestScanReadsVideos(AdminAPITestCase):
    """Videos saved by ComfyUI carry the graph in their comment tag."""

    def setUp(self):
        super().setUp()
        from utils import video_metadata

        video_metadata.ffprobe_path.cache_clear()
        self.addCleanup(video_metadata.ffprobe_path.cache_clear)

    @unittest.skipUnless(
        shutil.which("ffmpeg") and shutil.which("ffprobe"), "ffmpeg not installed"
    )
    async def test_scan_adds_prompts_found_in_videos(self):
        graph = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}}
        _write_mp4(
            self.output_dir / "clip.mp4", json.dumps({"prompt": json.dumps(graph)})
        )
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

        with self.assertNoLogs(level="ERROR"):
            body = await (
                await self.client.request("POST", "/prompt_manager/scan")
            ).text()

        self.assertIn('"added": 1', body)
        prompts = self.api.db.get_recent_prompts(limit=10)["prompts"]
        self.assertEqual([p["text"] for p in prompts], ["a cat"])
        images = self.api.db.get_prompt_images(prompts[0]["id"])
        self.assertTrue(images[0]["image_path"].endswith("clip.mp4"))

    async def test_scan_without_ffprobe_skips_videos_quietly(self):
        from unittest.mock import patch
        from utils import video_metadata

        (self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypisom")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

        with patch.object(video_metadata.shutil, "which", return_value=None):
            with self.assertNoLogs("prompt_manager", level="WARNING"):
                body = await (
                    await self.client.request("POST", "/prompt_manager/scan")
                ).text()

        self.assertIn('"processed": 1', body)
        self.assertIn('"found": 0', body)

    async def test_unreadable_image_is_a_warning_not_an_error(self):
        (self.output_dir / "bad.png").write_bytes(b"\x89PNG not really")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

        with self.assertLogs(level="WARNING") as cm:
            body = await (
                await self.client.request("POST", "/prompt_manager/scan")
            ).text()

        self.assertIn('"type": "complete"', body)
        self.assertTrue(any("bad.png" in line for line in cm.output))
        self.assertEqual([line for line in cm.output if line.startswith("ERROR")], [])


def _sse_payloads(body):
    return [
        json.loads(line[6:]) for line in body.splitlines() if line.startswith("data: ")
    ]


class TestDuplicateScanIsABackgroundJob(AdminAPITestCase):
    """The duplicate scan streams progress, survives the client, and is single-flight
    with the duplicate delete."""

    STREAM = "/prompt_manager/scan_duplicates/stream"
    STATUS = "/prompt_manager/scan_duplicates/status"

    def _two_dups_one_single(self):
        (self.output_dir / "a.png").write_bytes(b"same bytes")
        (self.output_dir / "b.png").write_bytes(b"same bytes")
        (self.output_dir / "c.png").write_bytes(b"other bytes")
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]

    def _gate_hashing(self):
        import threading

        gate = threading.Event()
        calls = []
        original = self.api._calculate_file_hash

        def gated(path):
            calls.append(path)
            gate.wait(5)
            return original(path)

        self.api._calculate_file_hash = gated
        return gate, calls

    async def test_stream_reports_progress_then_the_groups(self):
        self._two_dups_one_single()

        body = await (await self.client.request("POST", self.STREAM)).text()
        events = _sse_payloads(body)

        types = [e["type"] for e in events]
        self.assertIn("progress", types)
        self.assertEqual(types[-1], "complete")
        done = events[-1]
        self.assertEqual(done["processed"], 3)
        self.assertEqual(done["found"], 1)
        self.assertEqual(done["duplicates"][0]["count"], 2)
        self.assertEqual(
            sorted(i["filename"] for i in done["duplicates"][0]["images"]),
            ["a.png", "b.png"],
        )
        self.assertNotIn(self.tmpdir, body)

    async def test_second_request_attaches_and_delete_is_refused_meanwhile(self):
        self._two_dups_one_single()
        gate, calls = self._gate_hashing()

        first = asyncio.ensure_future(self.client.request("POST", self.STREAM))
        await asyncio.sleep(0.1)
        second = await asyncio.wait_for(self.client.request("POST", self.STREAM), 5)
        self.assertEqual(second.status, 200)
        delete = await self.client.request(
            "POST",
            "/prompt_manager/delete_duplicate_images",
            json={"image_paths": ["a.png"]},
        )
        self.assertEqual(delete.status, 409)
        running = await (await self.client.request("GET", self.STATUS)).json()
        self.assertTrue(running["running"])

        gate.set()
        first_body = await (await first).text()
        second_body = await second.text()
        finished = await (await self.client.request("GET", self.STATUS)).json()

        self.assertIn('"type": "complete"', first_body)
        self.assertIn('"type": "complete"', second_body)
        self.assertEqual(len(calls), 3)  # one scan, each file hashed once
        self.assertFalse(finished["running"])
        self.assertEqual(finished["last_event"]["found"], 1)

    async def test_status_is_idle_before_any_scan(self):
        data = await (await self.client.request("GET", self.STATUS)).json()
        self.assertEqual(data, {"success": True, "running": False, "last_event": None})

    async def test_failed_scan_streams_an_error_and_releases_the_lock(self):
        GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
        self.api._find_comfyui_output_dir = _raise

        body = await (await self.client.request("POST", self.STREAM)).text()
        self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
        legacy = await self.client.request("GET", "/prompt_manager/scan_duplicates")

        self.assertEqual(_sse_payloads(body)[-1]["type"], "error")
        self.assertNotIn("boom", body)
        self.assertEqual(legacy.status, 200)

    def test_hashing_runs_on_worker_threads_and_reports_progress(self):
        import threading
        import time

        for i in range(8):
            (self.output_dir / f"f{i}.png").write_bytes(f"bytes {i % 4}".encode())
        PromptManagerConfig.WORKER_THREADS = 4
        seen = set()
        lock = threading.Lock()
        original = self.api._calculate_file_hash

        def recording(path):
            with lock:
                seen.add(threading.current_thread().name)
            time.sleep(0.05)
            return original(path)

        self.api._calculate_file_hash = recording
        reports = []

        groups = self.api._find_duplicate_images_sync(
            str(self.output_dir),
            progress=lambda done, total: reports.append((done, total)),
        )

        self.assertEqual(len(groups), 4)
        self.assertGreaterEqual(len(seen), 2)
        self.assertEqual(reports[-1], (8, 8))
        self.assertEqual([t for _, t in reports], [8] * len(reports))
        self.assertEqual(sorted(d for d, _ in reports), [d for d, _ in reports])


if __name__ == "__main__":
    unittest.main()