jellyfin-scene-search/tests/test_routing.py

from __future__ import annotations

import unittest

from scene_search.routing import apply_context_routing


class CaptionRoutingTests(unittest.TestCase):
    def test_positive_frame_routes_nearby_frames_to_specialist(self) -> None:
        rows = [
            {"id": "a", "media_id": "movie", "timestamp_ms": 0},
            {"id": "b", "media_id": "movie", "timestamp_ms": 2000},
            {"id": "c", "media_id": "movie", "timestamp_ms": 8000},
            {"id": "d", "media_id": "other", "timestamp_ms": 2000},
        ]
        scores = {
            "a": 0.10,
            "b": 0.90,
            "c": 0.10,
            "d": 0.10,
        }

        decisions = apply_context_routing(
            rows,
            scores,
            threshold=0.50,
            context_ms=3000,
        )

        self.assertEqual(decisions["a"], ("specialist", "context"))
        self.assertEqual(decisions["b"], ("specialist", "score"))
        self.assertEqual(decisions["c"], ("general", "score"))
        self.assertEqual(decisions["d"], ("general", "score"))

    def test_threshold_is_inclusive(self) -> None:
        rows = [{"id": "a", "media_id": "movie", "timestamp_ms": 0}]
        scores = {"a": 0.50}

        decisions = apply_context_routing(
            rows,
            scores,
            threshold=0.50,
            context_ms=0,
        )

        self.assertEqual(decisions["a"], ("specialist", "score"))


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