Test shared search service and item grouping

2026-11-14 22:11:27 +00:00 · 2026-10-06 16:38:48 +01:00
parent e125b7ac33
commit d1e6ff5d91
1 changed files with 173 additions and 0 deletions
@@ -0,0 +1,173 @@
+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()