import unittest

from app.video_pipeline import VideoResultSmoother


def result(class_id, class_name, confidence, box):
    return {
        "bbox_xyxy": box,
        "cls_class_id": class_id,
        "cls_class_name": class_name,
        "cls_confidence": confidence,
        "show_cls_name": True,
    }


class VideoResultSmootherTests(unittest.TestCase):
    def test_one_low_confidence_label_does_not_flip_track(self):
        smoother = VideoResultSmoother(hold_frames=4, alpha=0.5)
        smoother.update([result(1, "白鹤", 0.95, [0, 0, 20, 20])], 0)
        second = smoother.update(
            [result(2, "白枕鹤", 0.20, [1, 1, 21, 21])],
            1,
        )
        self.assertEqual(second[0]["cls_class_name"], "白鹤")
        self.assertNotIn("_smoother_class_votes", second[0])

    def test_matching_does_not_require_same_predicted_class(self):
        smoother = VideoResultSmoother(hold_frames=4, alpha=0.5)
        smoother.update([result(1, "白鹤", 0.9, [0, 0, 20, 20])], 0)
        second = smoother.update(
            [result(2, "白枕鹤", 0.9, [2, 2, 22, 22])],
            1,
        )
        self.assertEqual(second[0]["bbox_xyxy"], [1.0, 1.0, 21.0, 21.0])


if __name__ == "__main__":
    unittest.main()