Optimize Python Tests (#24823)

* Optimize tests

* Address comments

* Fix test

frigate/test/__init__.py

@@ -0,0 +1,11 @@
+from frigate.config.config import stream_info_retriever
+
+
+def _unreachable_stream_info(ffmpeg, path: str) -> dict:
+    return {"width": 0, "height": 0, "fourcc": None}
+
+
+# Validating a camera without detect dimensions probes its stream. Test streams
+# are unreachable, so return the failed probe result instead of waiting ~12s per
+# stream for ffprobe to time out
+stream_info_retriever.get_stream_info = _unreachable_stream_info

frigate/test/__main__.py

@@ -0,0 +1,278 @@
+"""Run the unit test suite in parallel groups.
+
+Test modules are split into groups balanced by test count, similar to Home
+Assistant's split_tests.py, and each group runs in its own unittest process.
+
+Usage (from the repo root):
+    python3 -u -m frigate.test [-j JOBS] [--groups GROUPS] [--timeout SECONDS]
+"""
+
+import argparse
+import math
+import os
+import re
+import signal
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass, field
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+TEST_METHOD = re.compile(r"^\s+(?:async )?def test", re.MULTILINE)
+RAN_TESTS = re.compile(r"^Ran (\d+) tests? in", re.MULTILINE)
+
+# Each group pays for its own interpreter and imports, so smaller groups only
+# help when there are more cores than groups
+TESTS_PER_GROUP = 250
+
+# Each group runs in its own process group, which ctrl+c in the terminal does
+# not reach, so the runner passes the interrupt on to these
+interrupted = threading.Event()
+running_groups: set[int] = set()
+
+
+@dataclass
+class TestGroup:
+    index: int
+    modules: list[str] = field(default_factory=list)
+    weight: int = 0
+
+
+@dataclass
+class GroupResult:
+    group: TestGroup
+    returncode: int
+    output: str
+    duration: float
+
+    @property
+    def tests_run(self) -> int:
+        """Number of tests unittest reported running, 0 if it never finished."""
+        match = RAN_TESTS.search(self.output)
+        return int(match.group(1)) if match else 0
+
+
+def discover_modules() -> dict[str, int]:
+    """Find test modules the same way unittest discovery does.
+
+    Returns:
+        Module name mapped to the number of test methods it defines
+    """
+    modules: dict[str, int] = {}
+
+    for path in sorted((REPO_ROOT / "frigate").rglob("test*.py")):
+        relative = path.relative_to(REPO_ROOT)
+
+        # discovery only descends into packages
+        if not all(
+            (REPO_ROOT / parent / "__init__.py").exists()
+            for parent in relative.parents
+            if parent != Path(".")
+        ):
+            continue
+
+        name = ".".join(relative.with_suffix("").parts)
+        modules[name] = len(TEST_METHOD.findall(path.read_text()))
+
+    return modules
+
+
+def split_groups(modules: dict[str, int], count: int) -> list[TestGroup]:
+    """Split modules into groups with a similar number of tests.
+
+    Largest modules are placed first, each into the currently lightest group.
+    """
+    groups = [TestGroup(index) for index in range(1, count + 1)]
+
+    for name, weight in sorted(modules.items(), key=lambda m: m[1], reverse=True):
+        group = min(groups, key=lambda g: g.weight)
+        group.modules.append(name)
+        group.weight += weight
+
+    return [group for group in groups if group.modules]
+
+
+def signal_group(pgid: int, sig: signal.Signals) -> bool:
+    """Send a signal to every process in a process group.
+
+    Returns:
+        False if no process is left in the group
+    """
+    try:
+        os.killpg(pgid, sig)
+    except ProcessLookupError:
+        return False
+
+    return True
+
+
+def stop_group(pgid: int) -> None:
+    """Stop the processes a group's tests left running.
+
+    They get a moment to exit before being killed, so the multiprocessing
+    resource tracker, which ignores SIGTERM, can remove leaked shared memory
+    once the processes holding it have exited.
+    """
+    if signal_group(pgid, signal.SIGTERM):
+        time.sleep(2)
+        signal_group(pgid, signal.SIGKILL)
+
+
+def run_group(group: TestGroup, timeout: float) -> GroupResult:
+    """Run a group's modules in a unittest process of its own.
+
+    Args:
+        group: The modules to run
+        timeout: Seconds before the group is stopped
+
+    Returns:
+        The exit code and combined output of the unittest process
+    """
+    start = time.monotonic()
+
+    # a file instead of a pipe so a process left behind by a test can't keep
+    # the group from finishing by holding the pipe open
+    with tempfile.TemporaryFile("w+") as output:
+        proc = subprocess.Popen(
+            [sys.executable, "-u", "-X", "faulthandler", "-m", "unittest"]
+            + group.modules,
+            cwd=REPO_ROOT,
+            # a group in the background that read the terminal would be
+            # stopped, and a stopped process doesn't act on SIGABRT
+            stdin=subprocess.DEVNULL,
+            stdout=output,
+            stderr=subprocess.STDOUT,
+            # a process group of its own so the processes its tests start can
+            # be stopped with it
+            process_group=0,
+        )
+        running_groups.add(proc.pid)
+
+        # the runner may have been interrupted while this group was starting
+        if interrupted.is_set():
+            signal_group(proc.pid, signal.SIGINT)
+
+        timed_out = False
+
+        try:
+            proc.wait(timeout)
+        except subprocess.TimeoutExpired:
+            # faulthandler prints the stack of every thread on SIGABRT, which
+            # shows where the group is stuck
+            timed_out = True
+            proc.send_signal(signal.SIGABRT)
+            proc.wait()
+
+            # aborting skips the cleanup of the running test, which would
+            # have stopped the processes it started
+            stop_group(proc.pid)
+
+        running_groups.discard(proc.pid)
+        output.seek(0)
+        text = output.read()
+
+    if timed_out:
+        text += f"\nGroup timed out after {timeout:.0f}s\n"
+
+    return GroupResult(group, proc.returncode, text, time.monotonic() - start)
+
+
+def main() -> int:
+    """Split the tests into groups, run them in parallel, and report the results.
+
+    Returns:
+        The exit code, 1 if any group failed
+    """
+    cpus = os.cpu_count() or 1
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument(
+        "-j",
+        "--jobs",
+        type=int,
+        default=cpus,
+        help=f"number of groups to run at once (default: {cpus})",
+    )
+    parser.add_argument(
+        "--groups",
+        type=int,
+        help=(
+            "number of groups to split the tests into (default: one per job, "
+            f"or more to keep groups under about {TESTS_PER_GROUP} tests)"
+        ),
+    )
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        default=600,
+        help="seconds before a group is stopped and reported as failed",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="print the output of passing groups as well as failing ones",
+    )
+    args = parser.parse_args()
+
+    modules = discover_modules()
+    group_count = args.groups or max(
+        args.jobs, math.ceil(sum(modules.values()) / TESTS_PER_GROUP)
+    )
+    groups = split_groups(modules, group_count)
+    print(f"Running {len(groups)} test groups with {args.jobs} jobs", flush=True)
+
+    start = time.monotonic()
+    failed: list[GroupResult] = []
+    total_tests = 0
+
+    with ThreadPoolExecutor(max_workers=args.jobs) as executor:
+        futures = [executor.submit(run_group, group, args.timeout) for group in groups]
+
+        try:
+            for future in as_completed(futures):
+                result = future.result()
+                total_tests += result.tests_run
+                status = "OK" if result.returncode == 0 else "FAILED"
+                print(
+                    f"Group {result.group.index}/{len(groups)}: {status}, "
+                    f"{result.tests_run} tests in {result.duration:.1f}s",
+                    flush=True,
+                )
+
+                if result.returncode != 0:
+                    failed.append(result)
+
+                if result.returncode != 0 or args.verbose:
+                    print(result.output, flush=True)
+        except KeyboardInterrupt:
+            interrupted.set()
+            executor.shutdown(wait=False, cancel_futures=True)
+
+            for pgid in list(running_groups):
+                signal_group(pgid, signal.SIGINT)
+
+            raise
+
+    print(f"\nRan {total_tests} tests in {time.monotonic() - start:.1f}s")
+
+    if failed:
+        print(f"{len(failed)} of {len(groups)} groups FAILED, rerun one with:")
+
+        for result in sorted(failed, key=lambda r: r.group.index):
+            print(
+                f"  Group {result.group.index}: python3 -u -m unittest "
+                f"{' '.join(result.group.modules)}"
+            )
+
+        return 1
+
+    print("OK")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

frigate/test/const.py

@@ -1,4 +1,12 @@
 """Constants for testing."""
 
-TEST_DB = "test.db"
-TEST_DB_CLEANUPS = ["test.db", "test.db-shm", "test.db-wal"]
+import os
+import tempfile
+
+# Unique per process so test groups can run in parallel without sharing a database
+TEST_DB = os.path.join(tempfile.gettempdir(), f"frigate_test_{os.getpid()}.db")
+TEST_DB_CLEANUPS = [TEST_DB, f"{TEST_DB}-shm", f"{TEST_DB}-wal"]
+
+# The database is thrown away after each test, so skip syncing every commit to
+# disk, which is slow when several test processes are writing at once
+TEST_DB_PRAGMAS = {"synchronous": "off"}

frigate/test/http_api/base_http_test.py

@@ -18,7 +18,7 @@ from frigate.debug_replay import DebugReplayManager
 from frigate.jobs.export import JobStatePublisher
 from frigate.models import Event, Recordings, ReviewSegment
 from frigate.review.types import SeverityEnum
-from frigate.test.const import TEST_DB, TEST_DB_CLEANUPS
+from frigate.test.const import TEST_DB, TEST_DB_CLEANUPS, TEST_DB_PRAGMAS
 
 
 class AuthTestClient(TestClient):
@@ -38,12 +38,12 @@ class AuthTestClient(TestClient):
 class BaseTestHttp(unittest.TestCase):
     def setUp(self, models):
         # setup clean database for each test run
-        migrate_db = SqliteDatabase("test.db")
+        migrate_db = SqliteDatabase(TEST_DB, pragmas=TEST_DB_PRAGMAS)
         del logging.getLogger("peewee_migrate").handlers[:]
         router = Router(migrate_db)
         router.run()
         migrate_db.close()
-        self.db = SqliteQueueDatabase(TEST_DB)
+        self.db = SqliteQueueDatabase(TEST_DB, pragmas=TEST_DB_PRAGMAS)
         self.db.bind(models)
 
         # The export job manager broadcasts via JobStatePublisher on

frigate/test/http_api/test_http_latest_frame.py

@@ -1,11 +1,12 @@
 import os
 import shutil
-from unittest.mock import MagicMock
+import tempfile
+from unittest.mock import MagicMock, patch
 
 import cv2
 import numpy as np
 
-from frigate.output.preview import PREVIEW_CACHE_DIR, PREVIEW_FRAME_TYPE
+from frigate.output.preview import PREVIEW_FRAME_TYPE
 from frigate.test.http_api.base_http_test import AuthTestClient, BaseTestHttp
 
 
@@ -15,14 +16,15 @@ class TestHttpLatestFrame(BaseTestHttp):
         self.app = super().create_app()
         self.app.detected_frames_processor = MagicMock()
 
-        if os.path.exists(PREVIEW_CACHE_DIR):
-            shutil.rmtree(PREVIEW_CACHE_DIR)
-        os.makedirs(PREVIEW_CACHE_DIR)
-
-    def tearDown(self):
-        if os.path.exists(PREVIEW_CACHE_DIR):
-            shutil.rmtree(PREVIEW_CACHE_DIR)
-        super().tearDown()
+        # a directory of its own so test groups running in parallel can't
+        # delete each other's preview frames
+        self.preview_dir = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.preview_dir, ignore_errors=True)
+        cache_dir_patch = patch(
+            "frigate.output.preview.PREVIEW_CACHE_DIR", self.preview_dir
+        )
+        cache_dir_patch.start()
+        self.addCleanup(cache_dir_patch.stop)
 
     def test_latest_frame_fallback_to_preview(self):
         camera = "front_door"
@@ -45,7 +47,7 @@ class TestHttpLatestFrame(BaseTestHttp):
             2,
         )
         preview_path = os.path.join(
-            PREVIEW_CACHE_DIR, f"preview_{camera}-1234567890.0.{PREVIEW_FRAME_TYPE}"
+            self.preview_dir, f"preview_{camera}-1234567890.0.{PREVIEW_FRAME_TYPE}"
         )
         cv2.imwrite(preview_path, dummy_frame)
 
@@ -80,7 +82,7 @@ class TestHttpLatestFrame(BaseTestHttp):
 
         # 2. Create a dummy preview file
         preview_path = os.path.join(
-            PREVIEW_CACHE_DIR, f"preview_{camera}-999.0.{PREVIEW_FRAME_TYPE}"
+            self.preview_dir, f"preview_{camera}-999.0.{PREVIEW_FRAME_TYPE}"
         )
         cv2.imwrite(preview_path, dummy_frame)
 

frigate/test/test_export.py

@@ -617,9 +617,13 @@ class TestExportTimelineAlignment(unittest.TestCase):
         ):
             windows = exporter._staged_chapter_windows()
 
-        path = exporter._build_recording_segment_chapter_metadata_file(windows)
-        self.addCleanup(os.remove, path)
-        content = Path(path).read_text()
+        # the real cache dir only exists once the container has started
+        with tempfile.TemporaryDirectory() as tmpdir:
+            chapter_path = os.path.join(tmpdir, "chapters.txt")
+            exporter._chapter_metadata_path = lambda: chapter_path  # type: ignore[method-assign]
+            path = exporter._build_recording_segment_chapter_metadata_file(windows)
+            self.assertEqual(path, chapter_path)
+            content = Path(chapter_path).read_text()
 
         self.assertIn("START=0\nEND=21500", content)
         self.assertIn("START=21500\nEND=43000", content)

frigate/test/test_media_auth.py

@@ -21,7 +21,7 @@ from frigate.api.media_auth import (
 )
 from frigate.config import FrigateConfig
 from frigate.models import Event, Export, Recordings, ReviewSegment
-from frigate.test.const import TEST_DB, TEST_DB_CLEANUPS
+from frigate.test.const import TEST_DB, TEST_DB_CLEANUPS, TEST_DB_PRAGMAS
 
 _CONFIG = {
     "mqtt": {"host": "mqtt"},
@@ -222,11 +222,11 @@ class TestExportResolution(unittest.TestCase):
     """Export resolution requires a DB lookup."""
 
     def setUp(self):
-        migrate_db = SqliteDatabase("test.db")
+        migrate_db = SqliteDatabase(TEST_DB, pragmas=TEST_DB_PRAGMAS)
         del logging.getLogger("peewee_migrate").handlers[:]
         Router(migrate_db).run()
         migrate_db.close()
-        self.db = SqliteQueueDatabase(TEST_DB)
+        self.db = SqliteQueueDatabase(TEST_DB, pragmas=TEST_DB_PRAGMAS)
         self.db.bind([Event, ReviewSegment, Recordings, Export])
         self.config = FrigateConfig(**_CONFIG)
 

frigate/test/test_notice_registry.py

@@ -12,7 +12,7 @@ from playhouse.sqliteq import SqliteQueueDatabase
 from frigate.models import Notice, NoticeStats
 from frigate.notices.registry import NoticeRegistry
 from frigate.notices.types import NOTICE_KINDS, NoticeKind, NoticeSeverity, notice_id
-from frigate.test.const import TEST_DB, TEST_DB_CLEANUPS
+from frigate.test.const import TEST_DB, TEST_DB_CLEANUPS, TEST_DB_PRAGMAS
 
 # kinds that switch on the lifecycle knobs, so the tests do not depend on the
 # values the real catalog picks
@@ -49,12 +49,12 @@ class TestNoticeKinds(unittest.TestCase):
 
 class RegistryTestCase(unittest.TestCase):
     def setUp(self):
-        migrate_db = SqliteDatabase("test.db")
+        migrate_db = SqliteDatabase(TEST_DB, pragmas=TEST_DB_PRAGMAS)
         del logging.getLogger("peewee_migrate").handlers[:]
         router = Router(migrate_db)
         router.run()
         migrate_db.close()
-        self.db = SqliteQueueDatabase(TEST_DB)
+        self.db = SqliteQueueDatabase(TEST_DB, pragmas=TEST_DB_PRAGMAS)
         self.db.bind([Notice, NoticeStats])
         self.registry = NoticeRegistry()
         self.listener = MagicMock()

frigate/test/test_preview_loader.py

@@ -1,23 +1,23 @@
 import os
 import shutil
+import tempfile
 import unittest
+from unittest.mock import patch
 
-from frigate.output.preview import (
-    PREVIEW_CACHE_DIR,
-    PREVIEW_FRAME_TYPE,
-    get_most_recent_preview_frame,
-)
+from frigate.output.preview import PREVIEW_FRAME_TYPE, get_most_recent_preview_frame
 
 
 class TestPreviewLoader(unittest.TestCase):
     def setUp(self):
-        if os.path.exists(PREVIEW_CACHE_DIR):
-            shutil.rmtree(PREVIEW_CACHE_DIR)
-        os.makedirs(PREVIEW_CACHE_DIR)
-
-    def tearDown(self):
-        if os.path.exists(PREVIEW_CACHE_DIR):
-            shutil.rmtree(PREVIEW_CACHE_DIR)
+        # a directory of its own so test groups running in parallel can't
+        # delete each other's preview frames
+        self.preview_dir = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.preview_dir, ignore_errors=True)
+        cache_dir_patch = patch(
+            "frigate.output.preview.PREVIEW_CACHE_DIR", self.preview_dir
+        )
+        cache_dir_patch.start()
+        self.addCleanup(cache_dir_patch.stop)
 
     def test_get_most_recent_preview_frame_missing(self):
         self.assertIsNone(get_most_recent_preview_frame("test_camera"))
@@ -28,14 +28,14 @@ class TestPreviewLoader(unittest.TestCase):
         for ts in ["1000.0", "2000.0", "1500.0"]:
             with open(
                 os.path.join(
-                    PREVIEW_CACHE_DIR, f"preview_{camera}-{ts}.{PREVIEW_FRAME_TYPE}"
+                    self.preview_dir, f"preview_{camera}-{ts}.{PREVIEW_FRAME_TYPE}"
                 ),
                 "w",
             ) as f:
                 f.write(f"test_{ts}")
 
         expected_path = os.path.join(
-            PREVIEW_CACHE_DIR, f"preview_{camera}-2000.0.{PREVIEW_FRAME_TYPE}"
+            self.preview_dir, f"preview_{camera}-2000.0.{PREVIEW_FRAME_TYPE}"
         )
         self.assertEqual(get_most_recent_preview_frame(camera), expected_path)
 
@@ -45,7 +45,7 @@ class TestPreviewLoader(unittest.TestCase):
         for ts in ["1000.0", "2000.0"]:
             with open(
                 os.path.join(
-                    PREVIEW_CACHE_DIR, f"preview_{camera}-{ts}.{PREVIEW_FRAME_TYPE}"
+                    self.preview_dir, f"preview_{camera}-{ts}.{PREVIEW_FRAME_TYPE}"
                 ),
                 "w",
             ) as f:
@@ -53,7 +53,7 @@ class TestPreviewLoader(unittest.TestCase):
 
         # Test finding frame before or at 1500
         expected_path = os.path.join(
-            PREVIEW_CACHE_DIR, f"preview_{camera}-1000.0.{PREVIEW_FRAME_TYPE}"
+            self.preview_dir, f"preview_{camera}-1000.0.{PREVIEW_FRAME_TYPE}"
         )
         self.assertEqual(
             get_most_recent_preview_frame(camera, before=1500.0), expected_path
@@ -67,7 +67,7 @@ class TestPreviewLoader(unittest.TestCase):
         other_camera = "other_camera"
         with open(
             os.path.join(
-                PREVIEW_CACHE_DIR, f"preview_{other_camera}-3000.0.{PREVIEW_FRAME_TYPE}"
+                self.preview_dir, f"preview_{other_camera}-3000.0.{PREVIEW_FRAME_TYPE}"
             ),
             "w",
         ) as f:
@@ -78,12 +78,12 @@ class TestPreviewLoader(unittest.TestCase):
     def test_get_most_recent_preview_frame_hyphenated_camera(self):
         for name in ("preview_front-2000.0", "preview_front-door-3000.0"):
             with open(
-                os.path.join(PREVIEW_CACHE_DIR, f"{name}.{PREVIEW_FRAME_TYPE}"), "w"
+                os.path.join(self.preview_dir, f"{name}.{PREVIEW_FRAME_TYPE}"), "w"
             ) as f:
                 f.write("test")
 
         expected_path = os.path.join(
-            PREVIEW_CACHE_DIR, f"preview_front-2000.0.{PREVIEW_FRAME_TYPE}"
+            self.preview_dir, f"preview_front-2000.0.{PREVIEW_FRAME_TYPE}"
         )
         self.assertEqual(get_most_recent_preview_frame("front"), expected_path)
         self.assertEqual(
@@ -91,5 +91,5 @@ class TestPreviewLoader(unittest.TestCase):
         )
 
     def test_get_most_recent_preview_frame_no_directory(self):
-        shutil.rmtree(PREVIEW_CACHE_DIR)
+        shutil.rmtree(self.preview_dir)
         self.assertIsNone(get_most_recent_preview_frame("test_camera"))

frigate/test/test_storage.py

@@ -13,22 +13,28 @@ from frigate.config import FrigateConfig
 from frigate.const import STREAM_TYPE_MAIN, STREAM_TYPE_SUB
 from frigate.models import Event, Recordings
 from frigate.storage import MAX_CALCULATED_BANDWIDTH, StorageMaintainer
-from frigate.test.const import TEST_DB, TEST_DB_CLEANUPS
+from frigate.test.const import TEST_DB, TEST_DB_CLEANUPS, TEST_DB_PRAGMAS
 
 
 class TestHttp(unittest.TestCase):
     def setUp(self):
         # setup clean database for each test run
-        migrate_db = SqliteDatabase("test.db")
+        migrate_db = SqliteDatabase(TEST_DB, pragmas=TEST_DB_PRAGMAS)
         del logging.getLogger("peewee_migrate").handlers[:]
         router = Router(migrate_db)
         router.run()
         migrate_db.close()
-        self.db = SqliteQueueDatabase(TEST_DB)
+        self.db = SqliteQueueDatabase(TEST_DB, pragmas=TEST_DB_PRAGMAS)
         models = [Event, Recordings]
         self.db.bind(models)
         self.test_dir = tempfile.mkdtemp()
 
+        # The maintainers are never stopped, so a real subscriber leaks a zmq
+        # context whose garbage collection can hang a later test in term()
+        subscriber_patch = patch("frigate.storage.CameraConfigUpdateSubscriber")
+        subscriber_patch.start()
+        self.addCleanup(subscriber_patch.stop)
+
         self.minimal_config = {
             "mqtt": {"host": "mqtt"},
             "cameras": {