import json from types import SimpleNamespace from typing import Any, Callable, Iterator from unittest import mock import pytest from fastapi.testclient import TestClient from PIL import Image from immich_ml.config import log from immich_ml.main import app @pytest.fixture def pil_image() -> Image.Image: return Image.new("RGB", (600, 800)) @pytest.fixture def mock_get_model() -> Iterator[mock.Mock]: with mock.patch("immich_ml.models.cache.from_model_type", autospec=True) as mocked: yield mocked @pytest.fixture(scope="session") def deployed_app() -> Iterator[TestClient]: with TestClient(app) as client: yield client @pytest.fixture(scope="session") def responses() -> dict[str, Any]: responses: dict[str, Any] = json.load(open("responses.json", "r")) return responses @pytest.fixture(scope="session") def clip_model_cfg() -> dict[str, Any]: return { "embed_dim": 512, "vision_cfg": {"image_size": 224, "layers": 12, "width": 768, "patch_size": 32}, "text_cfg": {"context_length": 77, "vocab_size": 49408, "width": 512, "heads": 8, "layers": 12}, } @pytest.fixture(scope="session") def clip_preprocess_cfg() -> dict[str, Any]: return { "size": [224, 224], "mode": "RGB", "mean": [0.48145466, 0.4578275, 0.40821073], "std": [0.26862954, 0.26130258, 0.27577711], "interpolation": "bicubic", "resize_mode": "shortest", "fill_color": 0, } @pytest.fixture(scope="session") def clip_tokenizer_cfg() -> dict[str, Any]: return { "add_prefix_space": False, "added_tokens_decoder": { "49406": { "content": "<|startoftext|>", "lstrip": False, "normalized": True, "rstrip": False, "single_word": False, "special": True, }, "49407": { "content": "<|endoftext|>", "lstrip": False, "normalized": True, "rstrip": False, "single_word": False, "special": True, }, }, "bos_token": "<|startoftext|>", "clean_up_tokenization_spaces": True, "do_lower_case": True, "eos_token": "<|endoftext|>", "errors": "replace", "model_max_length": 77, "pad_token": "<|endoftext|>", "tokenizer_class": "CLIPTokenizer", "unk_token": "<|endoftext|>", } @pytest.fixture(scope="function") def providers(request: pytest.FixtureRequest) -> Iterator[mock.Mock]: marker = request.node.get_closest_marker("providers") if marker is None: raise ValueError("Missing marker 'providers'") providers = marker.args[0] with mock.patch("immich_ml.sessions.ort.ort.get_available_providers") as mocked: mocked.return_value = providers yield providers @pytest.fixture(scope="function") def ort_pybind() -> Iterator[mock.Mock]: with mock.patch("immich_ml.sessions.ort.ort.capi._pybind_state") as mocked: yield mocked @pytest.fixture(scope="function") def ov_device_ids(request: pytest.FixtureRequest, ort_pybind: mock.Mock) -> Iterator[mock.Mock]: marker = request.node.get_closest_marker("ov_device_ids") if marker is None: raise ValueError("Missing marker 'ov_device_ids'") ort_pybind.get_available_openvino_device_ids.return_value = marker.args[0] return ort_pybind @pytest.fixture(scope="function") def ort_session() -> Iterator[mock.Mock]: with mock.patch("immich_ml.sessions.ort.ort.InferenceSession") as mocked: yield mocked @pytest.fixture(scope="function") def stub_session() -> Callable[..., mock.Mock]: def _make(shape: tuple[Any, ...], outputs: Any = None, name: str = "input.1") -> mock.Mock: session = mock.Mock() session.get_inputs.return_value = [SimpleNamespace(name=name, shape=shape)] if outputs is not None: session.run.return_value = outputs return session return _make @pytest.fixture(scope="function") def ann_session() -> Iterator[mock.Mock]: with mock.patch("immich_ml.sessions.ann.Ann") as mocked: yield mocked @pytest.fixture(scope="function") def rknn_session() -> Iterator[mock.Mock]: with mock.patch("immich_ml.sessions.rknn.RknnPoolExecutor") as mocked: yield mocked @pytest.fixture(scope="function") def rmtree() -> Iterator[mock.Mock]: with mock.patch("immich_ml.models.base.rmtree", autospec=True) as mocked: mocked.avoids_symlink_attacks = True yield mocked @pytest.fixture(scope="function") def path() -> Iterator[mock.Mock]: path = mock.MagicMock() path.exists.return_value = True path.is_dir.return_value = True path.is_file.return_value = True path.with_suffix.return_value = path path.return_value = path with mock.patch("immich_ml.models.base.Path", return_value=path) as mocked: yield mocked @pytest.fixture(scope="function") def info() -> Iterator[mock.Mock]: with mock.patch.object(log, "info") as mocked: yield mocked @pytest.fixture(scope="function") def warning() -> Iterator[mock.Mock]: with mock.patch.object(log, "warning") as mocked: yield mocked @pytest.fixture(scope="function") def exception() -> Iterator[mock.Mock]: with mock.patch.object(log, "exception") as mocked: yield mocked @pytest.fixture(scope="function") def snapshot_download() -> Iterator[mock.Mock]: with mock.patch("immich_ml.models.base.snapshot_download") as mocked: yield mocked