ComfyAssets-ComfyUI_PromptManager/tests/test_job_admission.py

"""Registry admission classes and the legacy routes that hold the registry.

Record-changing work never runs together (Activity design): a ``records``
job is refused while any ``records`` or ``guarded`` job runs, a ``guarded``
job (thumbnails) while a ``records`` job runs, and ``read`` jobs only share
the one-per-kind rule. The legacy admin and gallery routes that change
records or files take a registry hold for their duration, so a workspace job
and a legacy-page operation refuse each other, naming the one that runs.

Fake runners and gated executor calls only; ordering is driven by events.
"""

import asyncio
import hashlib
import json
import os
import sys
import threading
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch

_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 py.api.jobs import (  # noqa: E402
    ADMISSIONS,
    Job,
    JobBusy,
    JobRegistry,
)

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


def _gated(gate):
    async def runner(job):
        await gate.wait()

    return runner


async def _noop(job):
    return None


class TestAdmissionClasses(unittest.IsolatedAsyncioTestCase):

    def _start(self, registry, kind, admission, gate=None, **kwargs):
        return registry.start(
            kind,
            kind.title(),
            admission=admission,
            scope={},
            runner=_gated(gate) if gate is not None else _noop,
            **kwargs,
        )

    async def test_admission_names_and_explicit_read(self):
        self.assertEqual(ADMISSIONS, ("read", "guarded", "records", "exclusive"))
        registry = JobRegistry()
        job = registry.start(
            "report", "Report", admission="read", scope={}, runner=_noop
        )
        await asyncio.wait_for(job.task, 2)
        self.assertEqual(job.admission, "read")
        self.assertFalse(job.record_changing)
        with self.assertRaises(ValueError):
            registry.start("x", "X", admission="bogus", scope={}, runner=_noop)
        with self.assertRaises(ValueError):
            registry.start(
                "x",
                "X",
                admission="read",
                record_changing=True,
                scope={},
                runner=_noop,
            )

    async def test_start_requires_an_admission_class(self):
        # A job that forgets its class must not silently run as "read" beside
        # a scan or restore: every start call decides.
        registry = JobRegistry()
        with self.assertRaises(TypeError) as raised:
            registry.start("review", "Review", scope={}, runner=_noop)
        self.assertIn("admission", str(raised.exception))
        self.assertEqual(registry.list(), [])
        job = registry.start(
            "review", "Review", record_changing=False, scope={}, runner=_noop
        )
        await asyncio.wait_for(job.task, 2)
        self.assertEqual(job.admission, "read")

    async def test_record_changing_is_an_alias_for_records(self):
        registry = JobRegistry()
        gate = asyncio.Event()
        job = registry.start(
            "scan", "Scan", record_changing=True, scope={}, runner=_gated(gate)
        )
        self.assertEqual(job.admission, "records")
        self.assertTrue(job.record_changing)
        self.assertEqual(job.summary()["admission"], "records")
        self.assertIs(
            registry.blocker("image-tags", record_changing=True), job
        )  # the existing call shape keeps working
        gate.set()
        await asyncio.wait_for(job.task, 2)

    async def test_records_job_refused_while_guarded_runs(self):
        registry = JobRegistry()
        gate = asyncio.Event()
        thumbs = self._start(registry, "thumbnails", "guarded", gate)
        with self.assertRaises(JobBusy) as raised:
            self._start(registry, "scan", "records")
        self.assertIs(raised.exception.running, thumbs)
        self.assertIs(registry.blocker("restore", admission="records"), thumbs)
        gate.set()
        await asyncio.wait_for(thumbs.task, 2)
        scan = self._start(registry, "scan", "records")
        await asyncio.wait_for(scan.task, 2)
        self.assertEqual(scan.state, "finished")

    async def test_guarded_refused_while_records_runs(self):
        registry = JobRegistry()
        gate = asyncio.Event()
        scan = self._start(registry, "scan", "records", gate)
        with self.assertRaises(JobBusy) as raised:
            self._start(registry, "thumbnails", "guarded")
        self.assertIs(raised.exception.running, scan)
        # A second guarded kind may run beside a guarded one, not beside records.
        gate.set()
        await asyncio.wait_for(scan.task, 2)
        hold = asyncio.Event()
        thumbs = self._start(registry, "thumbnails", "guarded", hold)
        backup = self._start(registry, "backup", "guarded")
        await asyncio.wait_for(backup.task, 2)
        self.assertEqual(backup.state, "finished")
        hold.set()
        await asyncio.wait_for(thumbs.task, 2)

    async def test_read_jobs_run_in_parallel(self):
        registry = JobRegistry()
        gate = asyncio.Event()
        scan = self._start(registry, "scan", "records", gate)
        health = self._start(registry, "health", "read", gate)
        duplicates = self._start(registry, "duplicates", "read", gate)
        self.assertTrue(all(j.running for j in (scan, health, duplicates)))
        with self.assertRaises(JobBusy) as raised:  # one job of a kind
            self._start(registry, "health", "read")
        self.assertIs(raised.exception.running, health)
        gate.set()
        await asyncio.wait_for(
            asyncio.gather(scan.task, health.task, duplicates.task), 2
        )

    async def test_non_cancellable_job_ignores_cancel(self):
        registry = JobRegistry()
        gate = asyncio.Event()
        job = self._start(registry, "restore", "records", gate, cancellable=False)
        self.assertIs(job.summary()["cancellable"], False)
        self.assertFalse(registry.cancel(job.id))
        self.assertFalse(job.cancelled)
        gate.set()
        await asyncio.wait_for(job.task, 2)
        self.assertEqual(job.state, "finished")

    async def test_hold_registers_a_running_non_cancellable_job(self):
        registry = JobRegistry()
        held = registry.hold(
            "restore", "Database restore (admin page)", admission="records"
        )
        async with held as job:
            self.assertIs(registry.active("restore"), job)
            self.assertEqual(job.state, "running")
            self.assertEqual(
                (job.summary()["cancellable"], job.summary()["admission"]),
                (False, "records"),
            )
            self.assertFalse(registry.cancel(job.id))
            with self.assertRaises(JobBusy) as raised:
                self._start(registry, "scan", "records")
            self.assertIn("Database restore (admin page)", str(raised.exception))
        self.assertEqual(job.state, "finished")
        self.assertIsNone(registry.active("restore"))
        self.assertIs(registry.list()[0], job)

    async def test_hold_refused_without_waiting(self):
        registry = JobRegistry()
        gate = asyncio.Event()
        scan = self._start(registry, "scan", "records", gate)
        with self.assertRaises(JobBusy) as raised:
            async with registry.hold("maintenance", "Maintenance", admission="records"):
                self.fail("a refused hold must not run its body")
        self.assertIs(raised.exception.running, scan)
        self.assertEqual([j.kind for j in registry.list()], ["scan"])
        gate.set()
        await asyncio.wait_for(scan.task, 2)

    async def test_hold_released_on_exception(self):
        registry = JobRegistry()
        with self.assertRaises(RuntimeError):
            async with registry.hold(
                "maintenance", "Maintenance", admission="records"
            ) as job:
                raise RuntimeError("boom")
        self.assertEqual(job.state, "failed")
        again = registry.hold("maintenance", "Maintenance", admission="records")
        async with again as second:
            second.finish("failed")  # the body may end it itself
        self.assertEqual(second.state, "failed")

    async def test_hold_admits_on_entry_not_at_the_call(self):
        # A hold created but never entered (an await or an exception between
        # the call and ``async with``) must not leave a running records job
        # behind that blocks all record work until a restart.
        registry = JobRegistry()
        held = registry.hold("restore", "Restore", admission="records")
        self.assertEqual(registry.list(), [])
        self.assertIsNone(registry.blocker("scan", admission="records"))
        del held
        async with registry.hold("restore", "Restore", admission="records") as job:
            self.assertIs(registry.active("restore"), job)
        self.assertEqual(job.state, "finished")

    def test_job_defaults_stay_read_and_cancellable(self):
        job = Job("duplicates", "Duplicate scan")
        self.assertEqual((job.admission, job.cancellable), ("read", True))


class _LegacyHoldCase(AdminAPITestCase):
    """The admin test app plus a running workspace job on demand."""

    def _run_job(self, kind="scan", title="Scan for metadata", admission="records"):
        gate = asyncio.Event()
        job = self.api._jobs.start(
            kind, title, admission=admission, scope={}, runner=_gated(gate)
        )
        self.addAsyncCleanup(self._release, gate, job)
        return gate, job

    @staticmethod
    async def _release(gate, job):
        gate.set()
        await asyncio.wait_for(job.task, 5)

    def _db_digest(self):
        path = Path(self.api.db.model.db_path)
        return {
            suffix: hashlib.sha256(Path(f"{path}{suffix}").read_bytes()).hexdigest()
            for suffix in ("", "-wal")
            if Path(f"{path}{suffix}").exists()
        }

    async def _busy(self, resp, title="Scan for metadata"):
        body = await resp.json()
        self.assertEqual(resp.status, 409)
        self.assertEqual(body, {"success": False, "error": f"{title} already running"})


class TestLegacyRoutesHoldTheRegistry(_LegacyHoldCase):

    async def test_hold_is_listed_as_running_job(self):
        held = self.api._jobs.hold(
            "restore", "Database restore (admin page)", admission="records"
        )
        async with held:
            resp = await self.client.request("GET", "/prompt_manager/jobs")
            jobs = (await resp.json())["jobs"]
        self.assertEqual(len(jobs), 1)
        job = jobs[0]
        self.assertEqual(
            (job["kind"], job["title"], job["state"]),
            ("restore", "Database restore (admin page)", "running"),
        )
        self.assertIs(job["cancellable"], False)
        self.assertEqual(job["admission"], "records")

    async def test_legacy_maintenance_and_cleanup_refused_while_scan_runs(self):
        self.api.db.save_prompt(text="keep me", prompt_hash="h1")
        before = self._db_digest()
        self._run_job()
        resp = await self.client.request(
            "POST",
            "/prompt_manager/maintenance",
            json={"operations": ["cleanup_duplicates", "vacuum"]},
        )
        await self._busy(resp)
        resp = await self.client.request("POST", "/prompt_manager/cleanup")
        await self._busy(resp)
        self.assertEqual(self._db_digest(), before)

    async def test_read_only_maintenance_runs_during_scan(self):
        # The legacy statistics and consistency checks only read: they hold
        # the registry as "read", so a long scan does not block them.
        self.api.db.save_prompt(text="keep me", prompt_hash="h1")
        _gate, scan = self._run_job()
        listed = []
        real = self.api._run_maintenance_operations

        def observed(operations):
            listed.extend(
                (job.kind, job.admission, job.state) for job in self.api._jobs.list()
            )
            return real(operations)

        self.api._run_maintenance_operations = observed
        resp = await self.client.request(
            "POST",
            "/prompt_manager/maintenance",
            json={
                "operations": [
                    "statistics",
                    "check_consistency",
                    "check_hash_duplicates",
                ]
            },
        )
        body = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertIs(body["success"], True)
        self.assertEqual(
            set(body["results"]),
            {"statistics", "check_consistency", "check_hash_duplicates"},
        )
        self.assertIn(("maintenance", "read", "running"), listed)
        self.assertEqual(scan.state, "running")
        maintenance = self.api._jobs.list()[-1]
        self.assertEqual(
            (maintenance.kind, maintenance.admission, maintenance.state),
            ("maintenance", "read", "finished"),
        )

    async def test_mixed_maintenance_selection_is_refused_during_scan(self):
        # One write operation among read-only ones makes the run "records".
        self.api.db.save_prompt(text="keep me", prompt_hash="h1")
        before = self._db_digest()
        self._run_job()
        resp = await self.client.request(
            "POST",
            "/prompt_manager/maintenance",
            json={"operations": ["statistics", "vacuum"]},
        )
        await self._busy(resp)
        self.assertEqual(self._db_digest(), before)

    async def test_default_maintenance_selection_is_refused_during_scan(self):
        # No "operations" key means the default list, which vacuums.
        self.api.db.save_prompt(text="keep me", prompt_hash="h1")
        before = self._db_digest()
        self._run_job()
        resp = await self.client.request("POST", "/prompt_manager/maintenance", json={})
        await self._busy(resp)
        self.assertEqual(self._db_digest(), before)
        self.assertEqual([j.kind for j in self.api._jobs.list()], ["scan"])

    async def test_legacy_maintenance_refused_while_image_tags_runs(self):
        self._run_job("image-tags", "Add tags to 3 images")
        resp = await self.client.request("POST", "/prompt_manager/maintenance")
        await self._busy(resp, "Add tags to 3 images")

    async def test_scan_refused_while_legacy_maintenance_holds(self):
        entered = threading.Event()
        release = threading.Event()
        real = self.api._run_maintenance_operations

        def gated(operations):
            entered.set()
            release.wait(10)
            return real(operations)

        self.api._run_maintenance_operations = gated
        loop = asyncio.get_running_loop()
        maintenance = asyncio.ensure_future(
            self.client.request(
                "POST", "/prompt_manager/maintenance", json={"operations": ["vacuum"]}
            )
        )
        self.assertTrue(await loop.run_in_executor(None, entered.wait, 10))
        resp = await self.client.request("POST", "/prompt_manager/jobs/scan")
        body = await resp.json()
        self.assertEqual((resp.status, body["code"]), (409, "busy"))
        self.assertEqual(body["running"]["title"], "Maintenance (admin page)")
        release.set()
        self.assertEqual((await maintenance).status, 200)
        listed = (
            await (await self.client.request("GET", "/prompt_manager/jobs")).json()
        )["jobs"]
        self.assertEqual(
            [(j["kind"], j["state"]) for j in listed], [("maintenance", "finished")]
        )

    async def test_legacy_duplicate_delete_refused_while_scan_runs(self):
        victim = self.output_dir / "dup.png"
        victim.write_bytes(b"x")
        self._run_job()
        resp = await self.client.request(
            "POST",
            "/prompt_manager/delete_duplicate_images",
            json={"image_paths": [str(victim)]},
        )
        await self._busy(resp)
        self.assertTrue(victim.exists())
        # The duplicate family lock was not taken by the refused request.
        self.assertIsNotNone(self.api._acquire_long_job(self.api.LONG_JOB_DUPLICATES))

    async def test_lora_scan_refused_while_scan_runs(self):
        from py.config import IntegrationConfig

        lm = Path(self.tmpdir) / "custom_nodes" / "ComfyUI-Lora-Manager"
        (lm / "py").mkdir(parents=True)
        self._run_job()
        imported = []
        with (
            patch.multiple(
                IntegrationConfig, LORA_MANAGER_ENABLED=True, LORA_MANAGER_PATH=str(lm)
            ),
            patch("py.lora_utils.resolve_lora_manager_path", return_value=str(lm)),
            patch.object(
                self.api, "_lora_scan_stream", side_effect=lambda *a: imported.append(a)
            ),
        ):
            resp = await self.client.request("POST", "/prompt_manager/lora/scan")
            text = await resp.text()
        self.assertEqual(resp.status, 200)  # the legacy page reads a stream
        events = [json.loads(line[6:]) for line in text.split("\n\n") if line]
        # The admin page's modal shows ``status`` and closes on ``type: error``.
        message = "Scan for metadata already running"
        self.assertEqual(
            events,
            [{"type": "error", "status": message, "progress": 0, "message": message}],
        )
        self.assertEqual(imported, [])

    async def test_legacy_autotag_refused_while_scan_runs(self):
        # Both pages read this route as a stream (a 409 reads as a lost
        # connection there), so a refusal by another job is one error frame:
        # the admin page keys on ``type``, the gallery on ``status``.
        self._run_job()
        resp = await self.client.request(
            "POST", "/prompt_manager/autotag/start", json={}
        )
        text = await resp.text()
        self.assertEqual(resp.status, 200)
        self.assertEqual(resp.headers["Content-Type"], "text/event-stream")
        events = [json.loads(line[6:]) for line in text.split("\n\n") if line]
        message = "Scan for metadata already running"
        self.assertEqual(
            events, [{"type": "error", "status": "error", "message": message}]
        )
        self.assertEqual(
            [j.kind for j in self.api._jobs.list() if j.kind == "autotag"], []
        )

    async def test_second_legacy_autotag_batch_keeps_its_json_409(self):
        from py.api.autotag_routes import _BUSY_ERROR

        self._run_job("autotag", "Auto-tag images (legacy page)")
        resp = await self.client.request(
            "POST", "/prompt_manager/autotag/start", json={}
        )
        self.assertEqual(resp.status, 409)
        self.assertEqual(await resp.json(), {"success": False, "error": _BUSY_ERROR})

    async def test_legacy_payloads_unchanged_when_idle(self):
        resp = await self.client.request(
            "POST", "/prompt_manager/maintenance", json={"operations": ["vacuum"]}
        )
        body = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertEqual(
            set(body),
            {"success", "operations_completed", "all_successful", "results", "message"},
        )
        resp = await self.client.request("POST", "/prompt_manager/cleanup")
        body = await resp.json()
        self.assertEqual(resp.status, 200)
        self.assertEqual(set(body), {"success", "message", "duplicates_removed"})
        resp = await self.client.request(
            "POST", "/prompt_manager/delete_duplicate_images", json={"image_paths": []}
        )
        self.assertEqual(
            (resp.status, await resp.json()),
            (400, {"success": False, "error": "No image paths provided"}),
        )
        # Each hold ended with the request; refused input never took one.
        listed = await self._listed()
        self.assertEqual(
            [(j["kind"], j["state"]) for j in listed],
            [("maintenance", "finished"), ("maintenance", "finished")],
        )
        self.assertEqual(
            [j["title"] for j in listed],
            ["Prompt cleanup (admin page)", "Maintenance (admin page)"],
        )

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


class TestHeldJobStates(_LegacyHoldCase):
    """Activity lists what a held legacy operation really did (Minor 7)."""

    async def _listed(self):
        resp = await self.client.request("GET", "/prompt_manager/jobs")
        return [(j["kind"], j["state"]) for j in (await resp.json())["jobs"]]

    async def test_refused_input_takes_no_hold(self):
        resp = await self.client.request(
            "POST", "/prompt_manager/delete_duplicate_images", json={}
        )
        self.assertEqual(resp.status, 400)
        resp = await self.client.request(
            "POST",
            "/prompt_manager/images/generate-thumbnails",
            data=b"{not json",
            headers={"Content-Type": "application/json"},
        )
        self.assertEqual(resp.status, 400)
        self.assertEqual(await self._listed(), [])

    async def test_thumbnail_stream_that_ends_in_error_is_failed(self):
        with patch.object(self.api, "_get_all_output_dirs", return_value=[]):
            resp = await self.client.request(
                "POST", "/prompt_manager/images/generate-thumbnails/progress"
            )
            text = await resp.text()
        self.assertEqual(resp.status, 200)
        self.assertIn("event: error", text)
        self.assertEqual(await self._listed(), [("thumbnails", "failed")])

    async def test_legacy_autotag_that_ends_in_error_is_failed(self):
        service = MagicMock()
        service.get_models_status.return_value = {}  # model not downloaded
        with patch("py.autotag.get_autotag_service", return_value=service):
            resp = await self.client.request(
                "POST", "/prompt_manager/autotag/start", json={}
            )
            text = await resp.text()
        self.assertEqual(resp.status, 200)
        self.assertIn('"type": "error"', text)
        self.assertEqual(await self._listed(), [("autotag", "failed")])


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