mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-04-10 09:07:37 +03:00
Refactor tests to extract duplicate logic into helper method
Co-authored-by: Teagan42 <2989925+Teagan42@users.noreply.github.com>
This commit is contained in:
parent
a2ad45e357
commit
12de5011fe
@ -6,6 +6,38 @@ from unittest.mock import MagicMock
|
|||||||
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"""
|
||||||
|
|
||||||
|
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):
|
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"""
|
"""Test that zones are included in sub_label classification messages when object is in zones"""
|
||||||
# Create a simple mock requestor
|
# Create a simple mock requestor
|
||||||
@ -18,18 +50,10 @@ class TestCustomObjectClassificationZones(unittest.TestCase):
|
|||||||
"current_zones": ["driveway", "front_yard"],
|
"current_zones": ["driveway", "front_yard"],
|
||||||
}
|
}
|
||||||
|
|
||||||
# Simulate what the processor does when publishing sub_label classification
|
# Build classification data using helper
|
||||||
classification_data = {
|
classification_data = self._build_classification_data(
|
||||||
"type": "classification",
|
obj_data, "sub_label", "person_walking"
|
||||||
"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"]
|
|
||||||
|
|
||||||
requestor.send_data("tracked_object_update", json.dumps(classification_data))
|
requestor.send_data("tracked_object_update", json.dumps(classification_data))
|
||||||
|
|
||||||
@ -65,18 +89,11 @@ class TestCustomObjectClassificationZones(unittest.TestCase):
|
|||||||
"current_zones": [],
|
"current_zones": [],
|
||||||
}
|
}
|
||||||
|
|
||||||
# Simulate what the processor does when publishing sub_label classification
|
# Build classification data using helper
|
||||||
classification_data = {
|
classification_data = self._build_classification_data(
|
||||||
"type": "classification",
|
obj_data, "sub_label", "person_running"
|
||||||
"id": obj_data["id"],
|
)
|
||||||
"camera": obj_data["camera"],
|
classification_data["score"] = 0.87
|
||||||
"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"]
|
|
||||||
|
|
||||||
requestor.send_data("tracked_object_update", json.dumps(classification_data))
|
requestor.send_data("tracked_object_update", json.dumps(classification_data))
|
||||||
|
|
||||||
@ -99,18 +116,12 @@ class TestCustomObjectClassificationZones(unittest.TestCase):
|
|||||||
"current_zones": ["site_entrance"],
|
"current_zones": ["site_entrance"],
|
||||||
}
|
}
|
||||||
|
|
||||||
# Simulate what the processor does when publishing attribute classification
|
# Build classification data using helper
|
||||||
classification_data = {
|
classification_data = self._build_classification_data(
|
||||||
"type": "classification",
|
obj_data, "attribute", "wearing_helmet"
|
||||||
"id": obj_data["id"],
|
)
|
||||||
"camera": obj_data["camera"],
|
classification_data["score"] = 0.92
|
||||||
"timestamp": 1234567890.0,
|
classification_data["model"] = "helmet_detector"
|
||||||
"model": "helmet_detector",
|
|
||||||
"attribute": "wearing_helmet",
|
|
||||||
"score": 0.92,
|
|
||||||
}
|
|
||||||
if obj_data.get("current_zones"):
|
|
||||||
classification_data["zones"] = obj_data["current_zones"]
|
|
||||||
|
|
||||||
requestor.send_data("tracked_object_update", json.dumps(classification_data))
|
requestor.send_data("tracked_object_update", json.dumps(classification_data))
|
||||||
|
|
||||||
@ -138,18 +149,12 @@ class TestCustomObjectClassificationZones(unittest.TestCase):
|
|||||||
"camera": "parking_lot",
|
"camera": "parking_lot",
|
||||||
}
|
}
|
||||||
|
|
||||||
# Simulate what the processor does when publishing attribute classification
|
# Build classification data using helper
|
||||||
classification_data = {
|
classification_data = self._build_classification_data(
|
||||||
"type": "classification",
|
obj_data, "attribute", "sedan"
|
||||||
"id": obj_data["id"],
|
)
|
||||||
"camera": obj_data["camera"],
|
classification_data["score"] = 0.95
|
||||||
"timestamp": 1234567890.0,
|
classification_data["model"] = "vehicle_type"
|
||||||
"model": "vehicle_type",
|
|
||||||
"attribute": "sedan",
|
|
||||||
"score": 0.95,
|
|
||||||
}
|
|
||||||
if obj_data.get("current_zones"):
|
|
||||||
classification_data["zones"] = obj_data["current_zones"]
|
|
||||||
|
|
||||||
requestor.send_data("tracked_object_update", json.dumps(classification_data))
|
requestor.send_data("tracked_object_update", json.dumps(classification_data))
|
||||||
|
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user