wehub-resource-sync
107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
"""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()
|