tests/unit/test_frame_annotator.py

from __future__ import annotations

from app.core.detector import Detection
from app.core.frame_annotator import MISSING_COLOR, PRESENT_COLOR, annotation_color


def make_detection(class_name: str) -> Detection:
    return Detection(
        class_id=0,
        class_name=class_name,
        confidence=0.9,
        x1=0,
        y1=0,
        x2=100,
        y2=100,
        color=(1, 2, 3),
    )


def test_ppe_present_classes_are_green():
    assert annotation_color(make_detection("Hardhat")) == PRESENT_COLOR
    assert annotation_color(make_detection("Mask")) == PRESENT_COLOR
    assert annotation_color(make_detection("Safety Vest")) == PRESENT_COLOR


def test_missing_ppe_classes_are_red():
    assert annotation_color(make_detection("NO-Hardhat")) == MISSING_COLOR
    assert annotation_color(make_detection("NO-Mask")) == MISSING_COLOR
    assert annotation_color(make_detection("NO-Safety Vest")) == MISSING_COLOR


def test_other_classes_keep_model_color():
    assert annotation_color(make_detection("Person")) == (1, 2, 3)