From 7d9e1e7a82a9b724312231be571c6270f6025eb5 Mon Sep 17 00:00:00 2001
From: Teagan Glenn <that@teagantotally.rocks>
Date: Sat, 31 Jan 2026 05:42:56 -0700
Subject: [PATCH] Store original mocks at on import

---
 frigate/test/test_custom_classification.py | 57 +++++++++++++---------
 1 file changed, 33 insertions(+), 24 deletions(-)

diff --git a/frigate/test/test_custom_classification.py b/frigate/test/test_custom_classification.py
index 9ec12ef9a..3a0644c48 100644
--- a/frigate/test/test_custom_classification.py
+++ b/frigate/test/test_custom_classification.py
@@ -3,6 +3,34 @@ import sys
 import unittest
 from unittest.mock import MagicMock
 
+MOCK_MODULES = [
+    "cv2",
+    "numpy",
+    "zmq",
+    "peewee",
+    "sherpa_onnx",
+    "pydantic",
+    "pydantic.fields",
+    "frigate.comms.inter_process",
+    "frigate.comms.event_metadata_updater",
+    "frigate.comms.embeddings_updater",
+    "frigate.log",
+    "frigate.const",
+    "frigate.util.builtin",
+    "frigate.util.object",
+    "tflite_runtime",
+    "tflite_runtime.interpreter",
+    "tensorflow",
+    "tensorflow.lite",
+    "tensorflow.lite.python",
+    "tensorflow.lite.python.interpreter",
+]
+
+ORIGINAL_MODULES = {
+    mod: sys.modules[mod]
+    for mod
+    in MOCK_MODULES
+}
 
 class TestCustomObjectClassificationZones(unittest.TestCase):
     """Test that zone information is correctly added to custom classification MQTT messages"""
@@ -185,36 +213,16 @@ class TestCustomObjectClassificationIntegration(unittest.TestCase):
 
     def setUp(self):
         """Import the processor after mocking dependencies"""
-        self.original_modules = sys.modules
-        # Mock all external dependencies before any imports
-        sys.modules["cv2"] = MagicMock()
-        sys.modules["numpy"] = MagicMock()
-        sys.modules["zmq"] = MagicMock()
-        sys.modules["peewee"] = MagicMock()
-        sys.modules["sherpa_onnx"] = MagicMock()
-
         # Create a better mock for pydantic to handle type annotations
         pydantic_mock = MagicMock()
         # Mock BaseModel as a simple class
         pydantic_mock.BaseModel = type("BaseModel", (), {})
         pydantic_mock.Field = MagicMock(return_value=None)
         pydantic_mock.ConfigDict = MagicMock(return_value={})
-        sys.modules["pydantic"] = pydantic_mock
-        sys.modules["pydantic.fields"] = MagicMock()
 
-        sys.modules["frigate.comms.inter_process"] = MagicMock()
-        sys.modules["frigate.comms.event_metadata_updater"] = MagicMock()
-        sys.modules["frigate.comms.embeddings_updater"] = MagicMock()
-        sys.modules["frigate.log"] = MagicMock()
-        sys.modules["frigate.const"] = MagicMock()
-        sys.modules["frigate.util.builtin"] = MagicMock()
-        sys.modules["frigate.util.object"] = MagicMock()
-        sys.modules["tflite_runtime"] = MagicMock()
-        sys.modules["tflite_runtime.interpreter"] = MagicMock()
-        sys.modules["tensorflow"] = MagicMock()
-        sys.modules["tensorflow.lite"] = MagicMock()
-        sys.modules["tensorflow.lite.python"] = MagicMock()
-        sys.modules["tensorflow.lite.python.interpreter"] = MagicMock()
+        for mod in MOCK_MODULES:
+            sys.modules[mod] = MagicMock()
+        sys.modules["pydantic"] = pydantic_mock
 
         # Import numpy after it's been mocked
         import numpy as np
@@ -232,7 +240,8 @@ class TestCustomObjectClassificationIntegration(unittest.TestCase):
             self.skipTest(f"Requires full Frigate environment: {e}")
 
     def tearDown(self):
-        sys.modules = self.original_modules
+        for mod in MOCK_MODULES:
+            sys.modules[mod] = ORIGINAL_MODULES[mode]
 
     def test_process_frame_with_zones_includes_zones_in_mqtt(self):
         """