initial commit

tests/test_pipeline_switches.py

@@ -0,0 +1,114 @@
+import unittest
+from unittest.mock import patch
+
+import numpy as np
+
+from app.config import settings
+from app.pipeline import BirdCascadePipeline
+
+
+class PipelineSwitchTests(unittest.TestCase):
+    setting_names = (
+        "enable_detection",
+        "enable_classification",
+        "enable_density",
+        "uniform_species_color",
+    )
+
+    def setUp(self):
+        self.original_values = {
+            name: getattr(settings, name)
+            for name in self.setting_names
+        }
+
+    def tearDown(self):
+        for name, value in self.original_values.items():
+            object.__setattr__(settings, name, value)
+
+    def test_all_models_can_be_disabled_without_loading_onnx(self):
+        object.__setattr__(settings, "enable_detection", False)
+        object.__setattr__(settings, "enable_classification", False)
+        object.__setattr__(settings, "enable_density", False)
+
+        with (
+            patch("app.pipeline.YOLOv5ONNXDetector") as detector_mock,
+            patch("app.pipeline.ConvNeXtONNXClassifier") as classifier_mock,
+            patch("app.pipeline.DensityEstimatorONNX") as density_mock,
+        ):
+            pipeline = BirdCascadePipeline()
+
+        self.assertIsNone(pipeline.detector)
+        self.assertIsNone(pipeline.classifier)
+        self.assertIsNone(pipeline.density_estimator)
+        detector_mock.assert_not_called()
+        classifier_mock.assert_not_called()
+        density_mock.assert_not_called()
+
+        output = pipeline.predict(
+            np.zeros((32, 32, 3), dtype=np.uint8),
+        )
+        self.assertEqual(output["results"], [])
+        self.assertIsNone(output["density"])
+
+    def test_detection_without_classification_keeps_detection_fields(self):
+        object.__setattr__(settings, "enable_detection", True)
+        object.__setattr__(settings, "enable_classification", False)
+        object.__setattr__(settings, "enable_density", False)
+
+        fake_detector_result = [
+            {
+                "bbox_xyxy": [4.0, 4.0, 20.0, 20.0],
+                "det_confidence": 0.9,
+                "det_class_id": 0,
+                "det_class_name": "bird",
+            }
+        ]
+
+        with (
+            patch(
+                "app.pipeline.read_classes",
+                return_value=["bird"],
+            ),
+            patch(
+                "app.pipeline.os.path.exists",
+                return_value=True,
+            ),
+            patch("app.pipeline.YOLOv5ONNXDetector") as detector_class,
+            patch("app.pipeline.ConvNeXtONNXClassifier") as classifier_class,
+        ):
+            detector_class.return_value.infer.return_value = fake_detector_result
+            pipeline = BirdCascadePipeline()
+            output = pipeline.predict(
+                np.zeros((32, 32, 3), dtype=np.uint8),
+            )
+
+        classifier_class.assert_not_called()
+        self.assertEqual(len(output["results"]), 1)
+        result = output["results"][0]
+        self.assertEqual(result["det_class_name"], "bird")
+        self.assertEqual(result["cls_class_id"], -1)
+        self.assertEqual(result["cls_class_name"], "")
+        self.assertEqual(result["cls_confidence"], 0.0)
+        self.assertFalse(result["show_cls_name"])
+
+    def test_uniform_species_color_uses_one_color_for_every_class(self):
+        pipeline = object.__new__(BirdCascadePipeline)
+
+        object.__setattr__(settings, "uniform_species_color", False)
+        first_color = pipeline._get_color_by_class_id(1)
+        second_color = pipeline._get_color_by_class_id(2)
+        self.assertNotEqual(first_color, second_color)
+
+        object.__setattr__(settings, "uniform_species_color", True)
+        self.assertEqual(
+            pipeline._get_color_by_class_id(1),
+            pipeline._get_color_by_class_id(2),
+        )
+        self.assertEqual(
+            pipeline._get_color_by_class_id(999),
+            (0, 190, 0),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()