Store original mocks at on import

This commit is contained in:
Teagan Glenn 2026-01-31 05:42:56 -07:00 committed by Teagan glenn
parent 44adc06f28
commit 7d9e1e7a82

View File

@ -3,6 +3,34 @@ import sys
import unittest import unittest
from unittest.mock import MagicMock 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): class TestCustomObjectClassificationZones(unittest.TestCase):
"""Test that zone information is correctly added to custom classification MQTT messages""" """Test that zone information is correctly added to custom classification MQTT messages"""
@ -185,36 +213,16 @@ class TestCustomObjectClassificationIntegration(unittest.TestCase):
def setUp(self): def setUp(self):
"""Import the processor after mocking dependencies""" """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 # Create a better mock for pydantic to handle type annotations
pydantic_mock = MagicMock() pydantic_mock = MagicMock()
# Mock BaseModel as a simple class # Mock BaseModel as a simple class
pydantic_mock.BaseModel = type("BaseModel", (), {}) pydantic_mock.BaseModel = type("BaseModel", (), {})
pydantic_mock.Field = MagicMock(return_value=None) pydantic_mock.Field = MagicMock(return_value=None)
pydantic_mock.ConfigDict = MagicMock(return_value={}) pydantic_mock.ConfigDict = MagicMock(return_value={})
sys.modules["pydantic"] = pydantic_mock
sys.modules["pydantic.fields"] = MagicMock()
sys.modules["frigate.comms.inter_process"] = MagicMock() for mod in MOCK_MODULES:
sys.modules["frigate.comms.event_metadata_updater"] = MagicMock() sys.modules[mod] = MagicMock()
sys.modules["frigate.comms.embeddings_updater"] = MagicMock() sys.modules["pydantic"] = pydantic_mock
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()
# Import numpy after it's been mocked # Import numpy after it's been mocked
import numpy as np import numpy as np
@ -232,7 +240,8 @@ class TestCustomObjectClassificationIntegration(unittest.TestCase):
self.skipTest(f"Requires full Frigate environment: {e}") self.skipTest(f"Requires full Frigate environment: {e}")
def tearDown(self): 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): def test_process_frame_with_zones_includes_zones_in_mqtt(self):
""" """