jellyfin-scene-search/tests/test_search_service.py

from __future__ import annotations

import tempfile
import unittest
from pathlib import Path
from unittest.mock import Mock

import numpy as np

from scene_search.database import connect, init_db, pack_vector
from scene_search.search_service import SearchService


class FakeEmbedder:
    def embed_text(self, text: str) -> np.ndarray:
        self.last_text = text
        return np.asarray([1.0, 0.0], dtype=np.float32)

    def match_probabilities(
        self,
        text_vector: np.ndarray,
        image_vectors: np.ndarray,
    ) -> np.ndarray:
        return np.asarray([0.9, 0.7, 0.6], dtype=np.float32)


class SearchServiceTests(unittest.TestCase):
    def test_search_returns_scenes_and_grouped_jellyfin_items(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            db_path = Path(tmp) / "scenes.sqlite3"
            init_db(db_path)

            with connect(db_path) as con:
                con.execute(
                    """
                    INSERT INTO media (
                        id, source_path, title, content_sha256, size_bytes, mtime_ns,
                        interval_ms, model_id, embedding_dim, source_kind,
                        jellyfin_item_id, jellyfin_library_id, jellyfin_item_type,
                        jellyfin_tags
                    )
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        "media-1",
                        "/media/example.mkv",
                        "Example Movie",
                        "hash",
                        1,
                        0,
                        10000,
                        "model",
                        2,
                        "jellyfin",
                        "jf-1",
                        "lib-1",
                        "Movie",
                        '["tag one"]',
                    ),
                )
                for number, timestamp in enumerate((0, 10000, 20000)):
                    con.execute(
                        """
                        INSERT INTO frames (
                            id, media_id, frame_number, timestamp_ms,
                            thumb_path, embedding_f16, caption
                        )
                        VALUES (?, ?, ?, ?, ?, ?, ?)
                        """,
                        (
                            f"frame-{number}",
                            "media-1",
                            number,
                            timestamp,
                            f"/tmp/frame-{number}.jpg",
                            pack_vector(np.asarray([1.0, 0.0], dtype=np.float32)),
                            f"caption {number}",
                        ),
                    )

            service = SearchService(
                db_path=db_path,
                meili_url="http://meili.test",
                meili_master_key="key",
                meili_index="scenes",
                model_id="model",
                device="cpu",
                semantic_ratio=0.8,
            )
            service._embedder = FakeEmbedder()  # type: ignore[assignment]
            service.meili = Mock()
            service.meili.search_vector.return_value = {
                "hits": [
                    {
                        "id": "frame-0",
                        "media_id": "media-1",
                        "jellyfin_item_id": "jf-1",
                        "jellyfin_library_id": "lib-1",
                        "jellyfin_item_type": "Movie",
                        "title": "Example Movie",
                        "timestamp_ms": 0,
                        "caption": "caption 0",
                        "jellyfin_tags": ["tag one"],
                        "_rankingScore": 0.95,
                    },
                    {
                        "id": "frame-1",
                        "media_id": "media-1",
                        "jellyfin_item_id": "jf-1",
                        "jellyfin_library_id": "lib-1",
                        "jellyfin_item_type": "Movie",
                        "title": "Example Movie",
                        "timestamp_ms": 10000,
                        "caption": "caption 1",
                        "jellyfin_tags": ["tag one"],
                        "_rankingScore": 0.85,
                    },
                    {
                        "id": "frame-2",
                        "media_id": "media-1",
                        "jellyfin_item_id": "jf-1",
                        "jellyfin_library_id": "lib-1",
                        "jellyfin_item_type": "Movie",
                        "title": "Example Movie",
                        "timestamp_ms": 20000,
                        "caption": "caption 2",
                        "jellyfin_tags": ["tag one"],
                        "_rankingScore": 0.75,
                    },
                ]
            }

            result = service.search("person outside", scene_limit=10, item_limit=5)

            self.assertEqual(result["scene_count"], 3)
            self.assertEqual(result["item_count"], 1)
            self.assertEqual(len(result["items"]), 1)
            self.assertEqual(result["items"][0]["jellyfin_item_id"], "jf-1")
            self.assertEqual(result["items"][0]["best_timestamp_ms"], 0)
            self.assertEqual(result["items"][0]["scene_matches"], 3)
            self.assertAlmostEqual(result["scenes"][0]["siglip_score"], 0.9, places=5)

            call = service.meili.search_vector.call_args
            self.assertEqual(call.args[1], "person outside")
            self.assertEqual(
                call.kwargs["filter_expression"],
                'source_kind = "jellyfin"',
            )

    def test_health_does_not_load_embedding_model(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            db_path = Path(tmp) / "scenes.sqlite3"
            service = SearchService(
                db_path=db_path,
                meili_url="http://meili.test",
                meili_master_key="key",
                meili_index="scenes",
                model_id="model",
                device="cpu",
                semantic_ratio=0.8,
            )
            service.meili = Mock()
            service.meili.stats.return_value = {"numberOfDocuments": 0}

            result = service.health()

            self.assertFalse(result["model_loaded"])
            self.assertEqual(result["status"], "ok")
            self.assertEqual(result["meili_documents"], 0)


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