From 12de5011fefe86bc192513bcc754898847ddd934 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sat, 31 Jan 2026 11:34:29 +0000
Subject: [PATCH] Refactor tests to extract duplicate logic into helper method

Co-authored-by: Teagan42 <2989925+Teagan42@users.noreply.github.com>
---
 frigate/test/test_custom_classification.py | 101 +++++++++++----------
 1 file changed, 53 insertions(+), 48 deletions(-)

diff --git a/frigate/test/test_custom_classification.py b/frigate/test/test_custom_classification.py
index 116d10694..978b88b92 100644
--- a/frigate/test/test_custom_classification.py
+++ b/frigate/test/test_custom_classification.py
@@ -6,6 +6,38 @@ from unittest.mock import MagicMock
 class TestCustomObjectClassificationZones(unittest.TestCase):
     """Test that zone information is correctly added to custom classification MQTT messages"""
 
+    def _build_classification_data(
+        self, obj_data, classification_type="sub_label", label="person_walking"
+    ):
+        """Helper method to build classification data with conditional zones.
+
+        Args:
+            obj_data: Object data dictionary containing id, camera, and optionally current_zones
+            classification_type: Either "sub_label" or "attribute"
+            label: The classification label
+
+        Returns:
+            Dictionary with classification data, including zones if applicable
+        """
+        classification_data = {
+            "type": "classification",
+            "id": obj_data["id"],
+            "camera": obj_data["camera"],
+            "timestamp": 1234567890.0,
+            "model": "test_classifier",
+            "score": 0.89,
+        }
+
+        if classification_type == "sub_label":
+            classification_data["sub_label"] = label
+        else:
+            classification_data["attribute"] = label
+
+        if obj_data.get("current_zones"):
+            classification_data["zones"] = obj_data["current_zones"]
+
+        return classification_data
+
     def test_sub_label_message_includes_zones_when_present(self):
         """Test that zones are included in sub_label classification messages when object is in zones"""
         # Create a simple mock requestor
@@ -18,18 +50,10 @@ class TestCustomObjectClassificationZones(unittest.TestCase):
             "current_zones": ["driveway", "front_yard"],
         }
 
-        # Simulate what the processor does when publishing sub_label classification
-        classification_data = {
-            "type": "classification",
-            "id": obj_data["id"],
-            "camera": obj_data["camera"],
-            "timestamp": 1234567890.0,
-            "model": "test_classifier",
-            "sub_label": "person_walking",
-            "score": 0.89,
-        }
-        if obj_data.get("current_zones"):
-            classification_data["zones"] = obj_data["current_zones"]
+        # Build classification data using helper
+        classification_data = self._build_classification_data(
+            obj_data, "sub_label", "person_walking"
+        )
 
         requestor.send_data("tracked_object_update", json.dumps(classification_data))
 
@@ -65,18 +89,11 @@ class TestCustomObjectClassificationZones(unittest.TestCase):
             "current_zones": [],
         }
 
-        # Simulate what the processor does when publishing sub_label classification
-        classification_data = {
-            "type": "classification",
-            "id": obj_data["id"],
-            "camera": obj_data["camera"],
-            "timestamp": 1234567890.0,
-            "model": "test_classifier",
-            "sub_label": "person_running",
-            "score": 0.87,
-        }
-        if obj_data.get("current_zones"):
-            classification_data["zones"] = obj_data["current_zones"]
+        # Build classification data using helper
+        classification_data = self._build_classification_data(
+            obj_data, "sub_label", "person_running"
+        )
+        classification_data["score"] = 0.87
 
         requestor.send_data("tracked_object_update", json.dumps(classification_data))
 
@@ -99,18 +116,12 @@ class TestCustomObjectClassificationZones(unittest.TestCase):
             "current_zones": ["site_entrance"],
         }
 
-        # Simulate what the processor does when publishing attribute classification
-        classification_data = {
-            "type": "classification",
-            "id": obj_data["id"],
-            "camera": obj_data["camera"],
-            "timestamp": 1234567890.0,
-            "model": "helmet_detector",
-            "attribute": "wearing_helmet",
-            "score": 0.92,
-        }
-        if obj_data.get("current_zones"):
-            classification_data["zones"] = obj_data["current_zones"]
+        # Build classification data using helper
+        classification_data = self._build_classification_data(
+            obj_data, "attribute", "wearing_helmet"
+        )
+        classification_data["score"] = 0.92
+        classification_data["model"] = "helmet_detector"
 
         requestor.send_data("tracked_object_update", json.dumps(classification_data))
 
@@ -138,18 +149,12 @@ class TestCustomObjectClassificationZones(unittest.TestCase):
             "camera": "parking_lot",
         }
 
-        # Simulate what the processor does when publishing attribute classification
-        classification_data = {
-            "type": "classification",
-            "id": obj_data["id"],
-            "camera": obj_data["camera"],
-            "timestamp": 1234567890.0,
-            "model": "vehicle_type",
-            "attribute": "sedan",
-            "score": 0.95,
-        }
-        if obj_data.get("current_zones"):
-            classification_data["zones"] = obj_data["current_zones"]
+        # Build classification data using helper
+        classification_data = self._build_classification_data(
+            obj_data, "attribute", "sedan"
+        )
+        classification_data["score"] = 0.95
+        classification_data["model"] = "vehicle_type"
 
         requestor.send_data("tracked_object_update", json.dumps(classification_data))