Aether/tests/services/test_model_cache_service.py

from types import SimpleNamespace
from typing import Any, Callable, cast

import pytest

from src.models.database import GlobalModel, Model
from src.services.cache.model_cache import ModelCacheService
from src.core.cache_service import CacheService


class _FakeQuery:
    def __init__(
        self,
        *,
        first_result: Any = None,
        all_result: list[Any] | None = None,
        on_all: Callable[[], None] | None = None,
    ) -> None:
        self._first_result = first_result
        self._all_result = all_result if all_result is not None else []
        self._on_all = on_all

    def join(self, *_args: object, **_kwargs: object) -> "_FakeQuery":
        return self

    def filter(self, *_args: object, **_kwargs: object) -> "_FakeQuery":
        return self

    def first(self) -> Any:
        return self._first_result

    def all(self) -> list[Any]:
        if self._on_all:
            self._on_all()
        return self._all_result


class _FakeSession:
    def __init__(self, *, direct_match: GlobalModel) -> None:
        self._direct_match = direct_match

    def query(self, *entities: object) -> "_FakeQuery":
        if entities == (GlobalModel,):
            return _FakeQuery(first_result=self._direct_match)

        # 如果 direct match 命中，不应再走 provider_model_name 分支
        if entities == (Model, GlobalModel):
            raise AssertionError(
                "provider_model_name query should not run when direct match exists"
            )

        raise AssertionError(f"Unexpected query entities: {entities}")


class _MappingIndexSession:
    def __init__(
        self,
        *,
        provider_mapping_rows: list[tuple[object, GlobalModel]],
    ) -> None:
        self._provider_mapping_rows = provider_mapping_rows
        self.provider_mapping_scan_count = 0
        self.model_global_query_count = 0

    def query(self, *entities: object) -> "_FakeQuery":
        if entities == (GlobalModel,):
            return _FakeQuery(first_result=None, all_result=[])

        if entities == (Model, GlobalModel):
            self.model_global_query_count += 1
            if self.model_global_query_count in {1, 3}:
                return _FakeQuery(all_result=[])
            if self.model_global_query_count == 2:
                return _FakeQuery(
                    all_result=self._provider_mapping_rows,
                    on_all=self._record_provider_mapping_scan,
                )
            raise AssertionError("provider_model_mappings 全量扫描被重复触发")

        raise AssertionError(f"Unexpected query entities: {entities}")

    def _record_provider_mapping_scan(self) -> None:
        self.provider_mapping_scan_count += 1


@pytest.mark.asyncio
async def test_resolve_global_model_prefers_direct_match(monkeypatch: pytest.MonkeyPatch) -> None:
    async def _fake_get(_key: str) -> None:
        return None

    async def _fake_set(_key: str, _value: object, ttl_seconds: int = 60) -> bool:  # noqa: ARG001
        return True

    monkeypatch.setattr(CacheService, "get", staticmethod(_fake_get))
    monkeypatch.setattr(CacheService, "set", staticmethod(_fake_set))

    global_model = GlobalModel(
        id="gm-1",
        name="claude-haiku-4-5-20251001",
        display_name="Claude Haiku 4.5",
        supported_capabilities=[],
        config={},
        default_tiered_pricing=None,
        default_price_per_request=None,
        is_active=True,
    )
    db = _FakeSession(direct_match=global_model)

    resolved = await ModelCacheService.resolve_global_model_by_name_or_mapping(
        cast(Any, db),
        cast(str, global_model.name),
    )
    assert resolved is global_model


@pytest.mark.asyncio
async def test_resolve_global_model_reuses_provider_mapping_index_cache(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    cache_store: dict[str, object] = {}

    async def _fake_get(key: str) -> object | None:
        return cache_store.get(key)

    async def _fake_set(key: str, value: object, ttl_seconds: int = 60) -> bool:  # noqa: ARG001
        cache_store[key] = value
        return True

    monkeypatch.setattr(CacheService, "get", staticmethod(_fake_get))
    monkeypatch.setattr(CacheService, "set", staticmethod(_fake_set))

    global_model_one = GlobalModel(
        id="gm-1",
        name="gpt-4o",
        display_name="GPT-4o",
        supported_capabilities=[],
        config={},
        default_tiered_pricing=None,
        default_price_per_request=None,
        is_active=True,
    )
    global_model_two = GlobalModel(
        id="gm-2",
        name="claude-3-7-sonnet",
        display_name="Claude 3.7 Sonnet",
        supported_capabilities=[],
        config={},
        default_tiered_pricing=None,
        default_price_per_request=None,
        is_active=True,
    )

    model_one = SimpleNamespace(
        id="m-1",
        provider_model_mappings=[{"name": "mapped-one"}],
    )
    model_two = SimpleNamespace(
        id="m-2",
        provider_model_mappings=[{"name": "mapped-two"}],
    )

    db = _MappingIndexSession(
        provider_mapping_rows=[
            (model_one, global_model_one),
            (model_two, global_model_two),
        ]
    )

    resolved_one = await ModelCacheService.resolve_global_model_by_name_or_mapping(
        cast(Any, db), "mapped-one"
    )
    resolved_two = await ModelCacheService.resolve_global_model_by_name_or_mapping(
        cast(Any, db), "mapped-two"
    )

    assert resolved_one is not None
    assert resolved_one.name == "gpt-4o"
    assert resolved_two is not None
    assert resolved_two.name == "claude-3-7-sonnet"
    assert db.provider_mapping_scan_count == 1


@pytest.mark.asyncio
async def test_invalidate_model_cache_clears_provider_mapping_index(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    deleted_keys: list[str] = []

    async def _fake_delete(key: str) -> bool:
        deleted_keys.append(key)
        return True

    monkeypatch.setattr(CacheService, "delete", staticmethod(_fake_delete))

    await ModelCacheService.invalidate_model_cache(
        model_id="model-1",
        provider_model_name="provider-model",
        provider_model_mappings=[{"name": "alias-model"}],
    )

    assert "model:id:model-1" in deleted_keys
    assert "global_model:resolve:provider-model" in deleted_keys
    assert "global_model:resolve:alias-model" in deleted_keys
    assert ModelCacheService.PROVIDER_MAPPING_INDEX_CACHE_KEY in deleted_keys