ComfyAssets-ComfyUI_PromptManager/tests/test_exit_window_threads.py

"""A thread or pool refused at interpreter shutdown is quiet (#375).

When the interpreter exits, ``threading._shutdown`` first shuts every
``ThreadPoolExecutor`` (``submit`` then raises "cannot schedule new futures
after interpreter shutdown") and joins the non-daemon threads; only then do
the ``atexit`` callbacks (the package teardown) run. Our daemon threads keep
working meanwhile. ``Thread.start`` is refused ("can't create new thread at
interpreter shutdown") only later, once the interpreter is finalizing; the
unit tests cover that case, the child process the pool case. Such a refusal
is logged once at DEBUG, with no traceback, and leaves nothing pending.

``threading.Thread`` is replaced only in the module under test (a copy of
the ``threading`` namespace), never process-wide. The child-process test runs
the real exit: a daemon thread queues previews while the main thread returns.
"""

import json
import logging
import os
import subprocess
import sys
import tempfile
import textwrap
import threading
import types
import unittest
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from unittest import mock

from py.api import thumbnail_auto
from utils import image_monitor, parallel

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
CHILD_CEILING_SECONDS = 60  # the hang watchdog allows 120 s per test
REPORT_PREFIX = "PM_EXIT_REPORT "
THREAD_REFUSAL = "can't create new thread at interpreter shutdown"
POOL_REFUSAL = "cannot schedule new futures after interpreter shutdown"


def _refusing_threading():
    """A ``threading`` namespace whose ``Thread.start`` is refused."""

    class RefusedThread(threading.Thread):
        def start(self):
            raise RuntimeError(THREAD_REFUSAL)

    namespace = types.SimpleNamespace(**vars(threading))
    namespace.Thread = RefusedThread
    return namespace


class _ShutPool(ThreadPoolExecutor):
    def submit(self, *args, **kwargs):
        raise RuntimeError(POOL_REFUSAL)


class _Capture(logging.Handler):
    def __init__(self):
        super().__init__(logging.DEBUG)
        self.records = []

    def emit(self, record):
        self.records.append(record)


def _capture(testcase, name):
    parallel._noted.clear()  # "once" is per process: start each test afresh
    logger = logging.getLogger(name)
    handler = _Capture()
    previous = logger.level
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)

    def restore():
        logger.removeHandler(handler)
        logger.setLevel(previous)

    testcase.addCleanup(restore)
    return handler.records


def _assert_one_debug_line(testcase, records):
    """Exactly one refusal line, at DEBUG and without a traceback."""
    refusals = [r for r in records if "interpreter is exiting" in r.getMessage()]
    testcase.assertEqual(len(refusals), 1, [r.getMessage() for r in records])
    testcase.assertEqual(refusals[0].levelno, logging.DEBUG)
    testcase.assertEqual([r for r in records if r.levelno > logging.DEBUG], [])
    testcase.assertTrue(all(r.exc_info is None for r in records))


class TestScheduleRefusedAtShutdown(unittest.TestCase):
    def setUp(self):
        self.handler = image_monitor.ImageGenerationHandler(
            mock.MagicMock(), mock.MagicMock()
        )
        self.addCleanup(self.handler.close)

    def test_schedule_returns_false_and_leaves_nothing_pending(self):
        records = _capture(self, "prompt_manager.image_monitor")
        with mock.patch.object(image_monitor, "threading", _refusing_threading()):
            self.assertFalse(self.handler.schedule("/out/a.png"))
            self.assertFalse(self.handler.schedule("/out/b.png"))
        self.assertEqual(self.handler.pending_count(), 0)
        _assert_one_debug_line(self, records)

    def test_another_start_failure_still_raises_and_leaves_nothing_pending(self):
        class Failing(threading.Thread):
            def start(self):
                raise RuntimeError("can't start new thread")

        namespace = types.SimpleNamespace(**vars(threading))
        namespace.Thread = Failing
        with mock.patch.object(image_monitor, "threading", namespace):
            with self.assertRaises(RuntimeError):
                self.handler.schedule("/out/a.png")
        self.assertEqual(self.handler.pending_count(), 0)

    def test_schedule_before_shutdown_is_unchanged(self):
        self.handler.processing_delay = 3600
        self.assertTrue(self.handler.schedule("/out/a.png"))
        self.assertEqual(self.handler.pending_count(), 1)


class TestPreviewQueueRefusedAtShutdown(unittest.TestCase):
    def thumbnailer(self, logger):
        made = thumbnail_auto.AutoThumbnailer(
            lambda *a, **k: {"action": "generated"}, lambda: True, logger
        )
        self.addCleanup(made.close)
        return made

    def test_enqueue_is_refused_quietly_and_drops_the_file(self):
        logger = logging.getLogger("prompt_manager.test_exit_window")
        records = _capture(self, logger.name)
        thumbnailer = self.thumbnailer(logger)
        src = Path("/out/a.png")
        with mock.patch.object(thumbnail_auto, "threading", _refusing_threading()):
            self.assertFalse(thumbnailer.enqueue(src, Path("/out/t/a.webp"), False))
            self.assertFalse(thumbnailer.enqueue(src, Path("/out/t/b.webp"), False))
        self.assertFalse(thumbnailer._worker_running)
        self.assertEqual(len(thumbnailer._queue), 0)
        self.assertEqual(thumbnailer._queued, set())
        self.assertTrue(thumbnailer.wait_idle(0))
        _assert_one_debug_line(self, records)


class TestMapParallelRefusedAtShutdown(unittest.TestCase):
    def test_a_shut_pool_runs_the_items_on_the_calling_thread(self):
        records = _capture(self, "prompt_manager.parallel")
        with mock.patch.object(parallel, "ThreadPoolExecutor", _ShutPool):
            self.assertEqual(
                parallel.map_parallel(lambda n: n * n, [1, 2, 3], 4), [1, 4, 9]
            )
        _assert_one_debug_line(self, records)

    def test_another_pool_error_still_raises(self):
        class Broken(ThreadPoolExecutor):
            def submit(self, *args, **kwargs):
                raise RuntimeError("something else")

        with mock.patch.object(parallel, "ThreadPoolExecutor", Broken):
            with self.assertRaises(RuntimeError):
                parallel.map_parallel(lambda n: n, [1, 2], 2)


CHILD = textwrap.dedent("""
    import threading

    # Registered before concurrent.futures is imported, so it runs after the
    # pool's own exit hook: from here on, a pool refuses new work.
    exiting = threading.Event()
    threading._register_atexit(exiting.set)

    import json
    import logging
    import os
    import sys
    from pathlib import Path

    root, tmp = sys.argv[1], sys.argv[2]
    sys.path.insert(0, root)

    from py.api.thumbnail_auto import AutoThumbnailer

    logging.basicConfig(level=logging.WARNING, stream=sys.stderr)
    logger = logging.getLogger("prompt_manager.exit_child")

    def generate(src, dst, is_video, *, stop, guard, track):
        Path(dst).write_bytes(b"preview")
        return {"action": "generated"}

    thumbnailer = AutoThumbnailer(
        generate, lambda: True, logger, workers=lambda: 4
    )
    queued = threading.Event()

    def late_monitor():
        exiting.wait()
        # One batch of two: the worker maps it on a pool.
        with thumbnailer._changed:
            for name in ("a", "b"):
                src = Path(tmp, name + ".png")
                src.write_bytes(b"x")
                thumbnailer.enqueue(src, Path(tmp, name + "_thumb.webp"), False)
        queued.set()

    def holder():
        # Non-daemon: the interpreter waits for it, inside the exit window.
        exiting.wait()
        queued.wait(30)
        idle = thumbnailer.wait_idle(30)
        written = sorted(
            n for n in os.listdir(tmp) if n.endswith("_thumb.webp")
        )
        print("PM_EXIT_REPORT " + json.dumps({"idle": idle, "written": written}))
        sys.stdout.flush()

    threading.Thread(target=late_monitor, daemon=True).start()
    threading.Thread(target=holder).start()
    # Return normally, as ComfyUI does after Ctrl+C.
    """)


class TestExitWindowChild(unittest.TestCase):
    def test_previews_queued_in_the_exit_window_log_no_traceback(self):
        with tempfile.TemporaryDirectory() as tmp:
            script = os.path.join(tmp, "child.py")
            with open(script, "w", encoding="utf-8") as handle:
                handle.write(CHILD)
            env = dict(os.environ)
            env.pop("PROMPT_MANAGER_CONFIG_PATH", None)
            env["PROMPT_MANAGER_DATA_DIR"] = tmp
            env["PYTHONIOENCODING"] = "utf-8"
            try:
                completed = subprocess.run(
                    [sys.executable, script, ROOT, tmp],
                    cwd=ROOT,
                    env=env,
                    capture_output=True,
                    timeout=CHILD_CEILING_SECONDS,
                    encoding="utf-8",
                    errors="replace",
                )
            except subprocess.TimeoutExpired as hung:
                self.fail(
                    f"the child did not exit within {CHILD_CEILING_SECONDS} s\n"
                    f"stdout:\n{hung.stdout}\nstderr:\n{hung.stderr}"
                )
        detail = f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}"
        self.assertEqual(completed.returncode, 0, detail)
        reports = [
            line[len(REPORT_PREFIX) :]
            for line in completed.stdout.splitlines()
            if line.startswith(REPORT_PREFIX)
        ]
        self.assertEqual(len(reports), 1, detail)
        report = json.loads(reports[0])
        self.assertTrue(report["idle"], detail)
        # The batch still ran: on the worker's own thread.
        self.assertEqual(report["written"], ["a_thumb.webp", "b_thumb.webp"], detail)
        self.assertNotIn("Traceback", completed.stderr, detail)


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