From 67ef6fa5282517db9dbdf92ccf40fc83f44485db Mon Sep 17 00:00:00 2001
From: Teagan Glenn <that@teagantotally.rocks>
Date: Sat, 31 Jan 2026 05:58:18 -0700
Subject: [PATCH] Mock modules but restore only modules that exist

---
 frigate/test/test_custom_classification.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/frigate/test/test_custom_classification.py b/frigate/test/test_custom_classification.py
index 1baabb2cb..89f6df010 100644
--- a/frigate/test/test_custom_classification.py
+++ b/frigate/test/test_custom_classification.py
@@ -26,7 +26,12 @@ MOCK_MODULES = [
     "tensorflow.lite.python.interpreter",
 ]
 
-ORIGINAL_MODULES = {mod: sys.modules[mod] for mod in MOCK_MODULES}
+ORIGINAL_MODULES = {
+    mod: sys.modules[mod]
+    for mod 
+    in MOCK_MODULES
+    if mod in sys.modules
+}
 
 
 class TestCustomObjectClassificationZones(unittest.TestCase):
@@ -238,7 +243,10 @@ class TestCustomObjectClassificationIntegration(unittest.TestCase):
 
     def tearDown(self):
         for mod in MOCK_MODULES:
-            sys.modules[mod] = ORIGINAL_MODULES[mod]
+            if mod in ORIGINAL_MODULES:
+                sys.modules[mod] = mod
+            else:
+                del sys.modules[mod]
 
     def test_process_frame_with_zones_includes_zones_in_mqtt(self):
         """