robertvoy-ComfyUI-Distributed/tests/test_image_list_lifecycle.py

"""Tests for periodic image-list server lifecycle cleanup."""
import asyncio
import time
import unittest
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import patch

from aiohttp import web

from image_list import lifecycle, runtime


class ImageListLifecycleTests(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        with runtime._LOCK:
            runtime._MASTER_JOBS.clear()
            runtime._WORKER_ITEMS.clear()

    def tearDown(self):
        with runtime._LOCK:
            runtime._MASTER_JOBS.clear()
            runtime._WORKER_ITEMS.clear()

    @staticmethod
    def _state(prompt_id=None, *, active=False, age=10_000.0):
        return SimpleNamespace(
            prompt_id=prompt_id,
            active=active,
            created_at=time.monotonic() - age,
        )

    @staticmethod
    def _server(queue):
        return SimpleNamespace(app=web.Application(), prompt_queue=queue)

    async def _start(self, server, *, interval=0.01, grace=0.01):
        lifecycle.install_lifecycle(
            server,
            runtime,
            interval_seconds=interval,
            orphan_grace_seconds=grace,
        )
        runner = web.AppRunner(server.app)
        await runner.setup()
        return runner

    async def test_unattended_sweep_removes_orphan_and_expired_staging(self):
        context_id = "orphan"
        with runtime._LOCK:
            cast(Any, runtime._MASTER_JOBS)[context_id] = self._state()
            runtime._WORKER_ITEMS["expired-token"] = (time.monotonic() - 1, {"image": "stale"})

        queue = SimpleNamespace(
            get_current_queue_volatile=lambda: ([], []),
        )
        server = self._server(queue)
        swept = asyncio.Event()
        sweep_worker_items = runtime._sweep_worker_items

        def track_worker_sweep():
            swept.set()
            sweep_worker_items()

        with patch.object(runtime, "_sweep_worker_items", side_effect=track_worker_sweep):
            runner = await self._start(server)
            try:
                await asyncio.wait_for(swept.wait(), timeout=0.5)
                self.assertNotIn(context_id, runtime._MASTER_JOBS)
                self.assertNotIn("expired-token", runtime._WORKER_ITEMS)
                with self.assertRaisesRegex(ValueError, "Unknown or expired"):
                    runtime.set_master_result(context_id, 0, "late-token", {})
            finally:
                await runner.cleanup()

    async def test_old_active_and_queued_contexts_survive_but_orphans_expire(self):
        states = {
            "active": self._state("active-prompt", active=True),
            "queued": self._state("queued-prompt"),
            "orphan": self._state("missing-prompt"),
            "unbound": self._state(None),
        }
        with runtime._LOCK:
            cast(Any, runtime._MASTER_JOBS).update(states)
        queue = SimpleNamespace(
            get_current_queue_volatile=lambda: (
                [[1, "active-prompt", {}]],
                [(2, "queued-prompt", {}, {}, None, None)],
            ),
        )
        runner = await self._start(self._server(queue))
        try:
            await asyncio.sleep(0.04)
            self.assertIs(runtime._MASTER_JOBS.get("active"), states["active"])
            self.assertIs(runtime._MASTER_JOBS.get("queued"), states["queued"])
            self.assertNotIn("orphan", runtime._MASTER_JOBS)
            self.assertNotIn("unbound", runtime._MASTER_JOBS)
        finally:
            await runner.cleanup()

    async def test_falls_back_to_nonvolatile_queue_reader(self):
        state = self._state("queued-prompt")
        with runtime._LOCK:
            cast(Any, runtime._MASTER_JOBS)["queued"] = state
        fallback_reads = 0

        def volatile_failure():
            raise RuntimeError("volatile view unavailable")

        def current_queue():
            nonlocal fallback_reads
            fallback_reads += 1
            return ([], [[1, "queued-prompt", {}]])

        queue = SimpleNamespace(
            get_current_queue_volatile=volatile_failure,
            get_current_queue=current_queue,
        )
        runner = await self._start(self._server(queue))
        try:
            await asyncio.sleep(0.04)
            self.assertGreaterEqual(fallback_reads, 1)
            self.assertIs(runtime._MASTER_JOBS.get("queued"), state)
        finally:
            await runner.cleanup()

    async def test_queue_read_error_protects_bound_contexts(self):
        with runtime._LOCK:
            cast(Any, runtime._MASTER_JOBS)["bound"] = self._state("bound-prompt")
            cast(Any, runtime._MASTER_JOBS)["unbound"] = self._state(None)

        def fail_read():
            raise RuntimeError("queue unavailable")

        queue = SimpleNamespace(
            get_current_queue_volatile=fail_read,
            get_current_queue=fail_read,
        )
        runner = await self._start(self._server(queue))
        try:
            await asyncio.sleep(0.04)
            self.assertIn("bound", runtime._MASTER_JOBS)
            self.assertNotIn("unbound", runtime._MASTER_JOBS)
        finally:
            await runner.cleanup()

    async def test_install_and_removal_are_idempotent_and_task_stops_on_cleanup(self):
        queue_reads = 0

        def read_queue():
            nonlocal queue_reads
            queue_reads += 1
            return ([], [])

        server = self._server(SimpleNamespace(get_current_queue_volatile=read_queue))
        lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01)
        lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01)
        self.assertEqual(len(server.app.cleanup_ctx), 1)

        with runtime._LOCK:
            cast(Any, runtime._MASTER_JOBS)["removed"] = self._state()
        runtime.remove_master_job("removed")
        runtime.remove_master_job("removed")
        self.assertNotIn("removed", runtime._MASTER_JOBS)

        runner = web.AppRunner(server.app)
        await runner.setup()
        await asyncio.sleep(0.035)
        self.assertGreaterEqual(queue_reads, 1)
        await runner.cleanup()
        reads_after_cleanup = queue_reads
        await asyncio.sleep(0.035)
        self.assertEqual(queue_reads, reads_after_cleanup)


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