immich/machine-learning/conftest.py

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