chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,106 @@
|
||||
"""Tests that disabled custom classification models are not registered or run."""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Mock TFLite before importing the maintainer / classification modules
|
||||
_MOCK_MODULES = [
|
||||
"tflite_runtime",
|
||||
"tflite_runtime.interpreter",
|
||||
"ai_edge_litert",
|
||||
"ai_edge_litert.interpreter",
|
||||
]
|
||||
for mod in _MOCK_MODULES:
|
||||
if mod not in sys.modules:
|
||||
sys.modules[mod] = MagicMock()
|
||||
|
||||
from frigate.data_processing.real_time.custom_classification import ( # noqa: E402
|
||||
CustomObjectClassificationProcessor,
|
||||
)
|
||||
from frigate.embeddings.maintainer import EmbeddingMaintainer # noqa: E402
|
||||
|
||||
|
||||
class TestCustomClassificationEnabledGating(unittest.TestCase):
|
||||
"""A model with enabled: false must not keep a processor registered."""
|
||||
|
||||
def _make_maintainer(self) -> EmbeddingMaintainer:
|
||||
# Bypass the heavy __init__; only the attributes touched by the
|
||||
# config update path are needed for these tests.
|
||||
maintainer = EmbeddingMaintainer.__new__(EmbeddingMaintainer)
|
||||
maintainer.realtime_processors = []
|
||||
maintainer.config = MagicMock()
|
||||
maintainer.config.classification.custom = {}
|
||||
maintainer.requestor = MagicMock()
|
||||
maintainer.metrics = MagicMock()
|
||||
maintainer.event_metadata_publisher = MagicMock()
|
||||
return maintainer
|
||||
|
||||
def _make_model_config(self, name: str, enabled: bool) -> MagicMock:
|
||||
model_config = MagicMock()
|
||||
model_config.name = name
|
||||
model_config.enabled = enabled
|
||||
model_config.state_config = None
|
||||
return model_config
|
||||
|
||||
def _make_processor(self, name: str) -> MagicMock:
|
||||
processor = MagicMock(spec=CustomObjectClassificationProcessor)
|
||||
processor.model_config = MagicMock()
|
||||
processor.model_config.name = name
|
||||
return processor
|
||||
|
||||
def test_disabled_update_tears_down_existing_processor(self):
|
||||
"""Toggling a running model to disabled shuts down and drops its processor."""
|
||||
maintainer = self._make_maintainer()
|
||||
processor = self._make_processor("atli")
|
||||
maintainer.realtime_processors = [processor]
|
||||
|
||||
maintainer._handle_custom_classification_update(
|
||||
"config/classification/custom/atli",
|
||||
self._make_model_config("atli", enabled=False),
|
||||
)
|
||||
|
||||
processor.shutdown.assert_called_once()
|
||||
self.assertEqual(maintainer.realtime_processors, [])
|
||||
|
||||
def test_disabled_update_does_not_register_processor(self):
|
||||
"""A disabled model that has no processor is never registered."""
|
||||
maintainer = self._make_maintainer()
|
||||
|
||||
maintainer._handle_custom_classification_update(
|
||||
"config/classification/custom/atli",
|
||||
self._make_model_config("atli", enabled=False),
|
||||
)
|
||||
|
||||
self.assertEqual(maintainer.realtime_processors, [])
|
||||
|
||||
def test_disabled_update_leaves_other_processors_untouched(self):
|
||||
"""Disabling one model must not affect other running processors."""
|
||||
maintainer = self._make_maintainer()
|
||||
other = self._make_processor("simbi")
|
||||
maintainer.realtime_processors = [other]
|
||||
|
||||
maintainer._handle_custom_classification_update(
|
||||
"config/classification/custom/atli",
|
||||
self._make_model_config("atli", enabled=False),
|
||||
)
|
||||
|
||||
other.shutdown.assert_not_called()
|
||||
self.assertEqual(maintainer.realtime_processors, [other])
|
||||
|
||||
def test_removed_model_tears_down_processor(self):
|
||||
"""A None payload (model deleted) still shuts down its processor."""
|
||||
maintainer = self._make_maintainer()
|
||||
processor = self._make_processor("atli")
|
||||
maintainer.realtime_processors = [processor]
|
||||
|
||||
maintainer._handle_custom_classification_update(
|
||||
"config/classification/custom/atli", None
|
||||
)
|
||||
|
||||
processor.shutdown.assert_called_once()
|
||||
self.assertEqual(maintainer.realtime_processors, [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user