mirror of
https://github.com/MrSleeps/jellyfin-scene-search.git
synced 2026-11-14 22:11:27 +00:00
Sleeps
67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import Mock
|
|
|
|
from scene_search.captioner import OllamaCaptioner, ensure_caption_prompt_file
|
|
|
|
|
|
class OllamaCaptionerTests(unittest.TestCase):
|
|
def test_caption_sends_image_and_returns_compact_text(self) -> None:
|
|
captioner = OllamaCaptioner("http://localhost:11434", "vision-test")
|
|
response = Mock()
|
|
response.ok = True
|
|
response.json.return_value = {
|
|
"message": {"content": " A person walking\n down a street. "}
|
|
}
|
|
captioner.session.post = Mock(return_value=response)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
image = Path(tmp) / "frame.jpg"
|
|
image.write_bytes(b"fake-jpeg")
|
|
result = captioner.caption(image, context_hint="Jellyfin tags: beach, holiday")
|
|
|
|
self.assertEqual(result, "A person walking down a street.")
|
|
call = captioner.session.post.call_args
|
|
payload = call.kwargs["json"]
|
|
self.assertEqual(payload["model"], "vision-test")
|
|
self.assertEqual(payload["messages"][0]["role"], "user")
|
|
self.assertIn("Jellyfin tags: beach, holiday", payload["messages"][0]["content"])
|
|
self.assertIn("actually visible", payload["messages"][0]["content"])
|
|
self.assertTrue(payload["messages"][0]["images"][0])
|
|
self.assertEqual(payload["options"]["temperature"], 0.0)
|
|
self.assertEqual(payload["options"]["num_predict"], 120)
|
|
|
|
def test_unload_releases_model_from_ollama(self) -> None:
|
|
captioner = OllamaCaptioner("http://localhost:11434", "vision-test")
|
|
response = Mock()
|
|
response.ok = True
|
|
captioner.session.post = Mock(return_value=response)
|
|
|
|
captioner.unload()
|
|
|
|
call = captioner.session.post.call_args
|
|
self.assertTrue(call.args[0].endswith("/api/generate"))
|
|
self.assertEqual(
|
|
call.kwargs["json"],
|
|
{"model": "vision-test", "keep_alive": 0},
|
|
)
|
|
|
|
def test_prompt_file_is_created_once_and_remains_editable(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
prompt_path = Path(tmp) / "caption_prompt.txt"
|
|
first = ensure_caption_prompt_file(prompt_path)
|
|
self.assertTrue(prompt_path.is_file())
|
|
self.assertTrue(first)
|
|
|
|
prompt_path.write_text("My test prompt\n", encoding="utf-8")
|
|
second = ensure_caption_prompt_file(prompt_path)
|
|
|
|
self.assertEqual(second, "My test prompt")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|