ComfyAssets-ComfyUI_PromptManager/tests/run_suite.py

"""Run the Python suite with a hang watchdog: ``python -m tests.run_suite``.

Why: on PR #349 a macOS job stopped inside the test suite and sat silent until
GitHub cancelled it after six hours, with no stack to read. This runner does
the same discovery as ``python -m unittest discover -s tests`` (same
``sys.path`` when started with ``-m`` from the repo root, same exit codes,
including 5 when no tests ran), and adds a watchdog built on
``faulthandler.dump_traceback_later``: when the timer is not re-armed in time,
every thread's stack goes to stderr and the process exits non-zero. It works
the same on Linux, macOS and Windows.

The timer is armed before discovery, re-armed at the start of every test with
PM_TEST_HANG_LIMIT seconds (default 120), and re-armed when the run ends with
PM_TEST_SHUTDOWN_LIMIT seconds (default 120) to cover module teardown and
interpreter exit. It is never cancelled. 120 s stays below asyncio's 300 s
wait for default-executor threads on Python 3.12 and 3.13, so a hang in an
``IsolatedAsyncioTestCase`` loop shutdown is dumped inside the test that hung.

Reading a dump: the last ``start: <test id>`` line names the test. If its
"ok" line comes after it, the hang is in that test's teardown or event-loop
shutdown; a dump after the summary is a hang at interpreter exit (a
non-daemon thread or an atexit handler).

A limit is a measurement, not a dial: change it only with new measurements
written next to it (TESTING.md "Time limits and hangs").
"""

import argparse
import faulthandler
import os
import sys
import unittest
from pathlib import Path

DEFAULT_PER_TEST_LIMIT = 120
DEFAULT_SHUTDOWN_LIMIT = 120
NO_TESTS_EXIT = 5
USAGE_EXIT = 2


class LimitError(ValueError):
    """An environment limit is not a positive integer."""


def read_limit(name, default):
    """Return the positive integer in env var ``name``, or ``default``."""
    raw = os.environ.get(name)
    if raw is None:
        return default
    text = raw.strip()
    if not text.isdigit() or int(text) <= 0:
        raise LimitError(f"run_suite: {name} must be a positive integer, got {raw!r}")
    return int(text)


def make_result_class(stream, per_test_limit, shutdown_limit):
    """Build a TextTestResult that names each test and re-arms the watchdog."""

    def arm(seconds):
        faulthandler.dump_traceback_later(seconds, exit=True, file=stream)

    class WatchdogResult(unittest.TextTestResult):
        def startTest(self, test):
            stream.write(f"start: {test.id()}\n")
            stream.flush()
            super().startTest(test)
            arm(per_test_limit)

        def stopTestRun(self):
            super().stopTestRun()
            arm(shutdown_limit)

    return WatchdogResult


def parse_args(argv):
    parser = argparse.ArgumentParser(
        prog="python -m tests.run_suite",
        description="Discover and run the tests with a hang watchdog.",
    )
    # As in unittest, -v and -q share one setting and the last one wins.
    parser.set_defaults(verbosity=1)
    parser.add_argument(
        "-v",
        "--verbose",
        dest="verbosity",
        action="store_const",
        const=2,
        help="verbose output",
    )
    parser.add_argument(
        "-q",
        "--quiet",
        dest="verbosity",
        action="store_const",
        const=0,
        help="quiet output",
    )
    parser.add_argument(
        "-s",
        "--start-directory",
        default=str(Path(__file__).resolve().parent),
        help="directory to start discovery (default: this tests folder)",
    )
    parser.add_argument(
        "-p", "--pattern", default="test*.py", help="pattern to match test files"
    )
    return parser.parse_args(argv)


def main(argv=None):
    args = parse_args(sys.argv[1:] if argv is None else argv)
    # Keep the real stderr: a test may replace sys.stderr, and the watchdog
    # needs a stream with a file descriptor.
    stream = sys.stderr
    try:
        per_test_limit = read_limit("PM_TEST_HANG_LIMIT", DEFAULT_PER_TEST_LIMIT)
        shutdown_limit = read_limit("PM_TEST_SHUTDOWN_LIMIT", DEFAULT_SHUTDOWN_LIMIT)
    except LimitError as error:
        stream.write(f"{error}\n")
        return USAGE_EXIT

    faulthandler.dump_traceback_later(per_test_limit, exit=True, file=stream)
    start = os.path.abspath(args.start_directory)
    suite = unittest.TestLoader().discover(start, args.pattern, top_level_dir=start)
    runner = unittest.TextTestRunner(
        stream=stream,
        verbosity=args.verbosity,
        resultclass=make_result_class(stream, per_test_limit, shutdown_limit),
    )
    result = runner.run(suite)
    if result.testsRun == 0 and not result.skipped:
        return NO_TESTS_EXIT
    return 0 if result.wasSuccessful() else 1


if __name__ == "__main__":
    sys.exit(main())