Aether/tests/services/test_provider_keys_key_command_service.py

from __future__ import annotations

import sys
import types
from contextlib import contextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Any, cast

import pytest

from src.core.exceptions import InvalidRequestException
from src.models.endpoint_models import EndpointAPIKeyCreate, EndpointAPIKeyUpdate


async def _noop_invalidate_models_list_cache() -> None:
    return None


_fake_models_service_module = types.ModuleType("src.api.base.models_service")
setattr(
    _fake_models_service_module, "invalidate_models_list_cache", _noop_invalidate_models_list_cache
)
sys.modules.setdefault("src.api.base.models_service", _fake_models_service_module)

from src.services.provider_keys import key_command_service as command_module
from src.services.provider_keys import key_side_effects as side_effects_module


class _NoQueryDB:
    def query(self, *args: Any, **kwargs: Any) -> Any:  # pragma: no cover - 防御断言
        _ = args, kwargs
        raise AssertionError("unexpected query call")


class _FakeQuery:
    def __init__(self, key: Any) -> None:
        self._key = key

    def filter(self, *args: Any, **kwargs: Any) -> _FakeQuery:
        _ = args, kwargs
        return self

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


class _FakeClearOAuthDB:
    def __init__(self, key: Any) -> None:
        self._key = key
        self.commit_count = 0

    def query(self, model: Any) -> _FakeQuery:
        _ = model
        return _FakeQuery(self._key)

    def commit(self) -> None:
        self.commit_count += 1


class _FakeQueryAll:
    def __init__(self, rows: list[Any]) -> None:
        self._rows = rows

    def filter(self, *args: Any, **kwargs: Any) -> _FakeQueryAll:
        _ = args, kwargs
        return self

    def all(self) -> list[Any]:
        return self._rows


def _build_key(**overrides: Any) -> SimpleNamespace:
    base: dict[str, Any] = {
        "auto_fetch_models": False,
        "allowed_models": None,
        "model_include_patterns": None,
        "model_exclude_patterns": None,
        "provider_id": "provider-1",
        "auth_type": "api_key",
        "api_formats": [],
        "provider": None,
    }
    base.update(overrides)
    return SimpleNamespace(**base)


def test_key_create_accepts_zero_max_probe_interval_minutes() -> None:
    payload = EndpointAPIKeyCreate.model_validate(
        {
            "name": "key-1",
            "api_key": "secret",
            "api_formats": ["openai:chat"],
            "max_probe_interval_minutes": 0,
        }
    )

    assert payload.max_probe_interval_minutes == 0


def test_key_update_accepts_zero_max_probe_interval_minutes() -> None:
    payload = EndpointAPIKeyUpdate.model_validate({"max_probe_interval_minutes": 0})

    assert payload.max_probe_interval_minutes == 0


def test_prepare_update_payload_auth_type_null_ignored() -> None:
    key = _build_key(auth_type="api_key")
    key_data = EndpointAPIKeyUpdate.model_validate({"auth_type": None})

    prepared = command_module._prepare_update_key_payload(
        db=cast(Any, _NoQueryDB()),
        key=cast(Any, key),
        key_id="key-1",
        key_data=key_data,
    )

    assert "auth_type" not in prepared.update_data


def test_prepare_update_payload_rejects_empty_api_key() -> None:
    key = _build_key(auth_type="api_key")
    key_data = EndpointAPIKeyUpdate.model_validate({"api_key": "   "})

    with pytest.raises(InvalidRequestException, match="api_key 不能为空"):
        command_module._prepare_update_key_payload(
            db=cast(Any, _NoQueryDB()),
            key=cast(Any, key),
            key_id="key-1",
            key_data=key_data,
        )


def test_prepare_update_payload_encrypts_empty_auth_config_dict(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    key = _build_key(auth_type="vertex_ai")
    key_data = EndpointAPIKeyUpdate.model_validate({"auth_config": {}})

    monkeypatch.setattr(command_module.crypto_service, "encrypt", lambda raw: f"ENC:{raw}")

    prepared = command_module._prepare_update_key_payload(
        db=cast(Any, _NoQueryDB()),
        key=cast(Any, key),
        key_id="key-1",
        key_data=key_data,
    )

    assert prepared.update_data["auth_config"] == "ENC:{}"


def test_prepare_update_payload_vertex_to_oauth_clears_auth_config(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    key = _build_key(auth_type="vertex_ai")
    key_data = EndpointAPIKeyUpdate.model_validate({"auth_type": "oauth"})

    monkeypatch.setattr(command_module.crypto_service, "encrypt", lambda raw: f"ENC:{raw}")

    prepared = command_module._prepare_update_key_payload(
        db=cast(Any, _NoQueryDB()),
        key=cast(Any, key),
        key_id="key-1",
        key_data=key_data,
    )

    assert prepared.update_data["auth_config"] is None
    assert prepared.update_data["api_key"] == "ENC:__placeholder__"


def test_validate_vertex_api_formats_api_key_allows_gemini_only() -> None:
    command_module._validate_vertex_api_formats("vertex_ai", "api_key", ["gemini:chat"])

    with pytest.raises(InvalidRequestException, match="claude:chat"):
        command_module._validate_vertex_api_formats("vertex_ai", "api_key", ["claude:chat"])


def test_validate_vertex_api_formats_service_account_allows_gemini_and_claude() -> None:
    command_module._validate_vertex_api_formats("vertex_ai", "service_account", ["claude:chat"])
    command_module._validate_vertex_api_formats("vertex_ai", "service_account", ["gemini:chat"])
    command_module._validate_vertex_api_formats(
        "vertex_ai", "service_account", ["gemini:chat", "claude:chat"]
    )

    with pytest.raises(InvalidRequestException, match="openai:chat"):
        command_module._validate_vertex_api_formats("vertex_ai", "service_account", ["openai:chat"])


def test_prepare_update_payload_allows_unrelated_update_for_legacy_vertex_combo() -> None:
    key = _build_key(
        auth_type="service_account",
        api_formats=["gemini:chat"],
        provider=SimpleNamespace(provider_type="vertex_ai"),
    )
    key_data = EndpointAPIKeyUpdate.model_validate({"name": "legacy-key"})

    prepared = command_module._prepare_update_key_payload(
        db=cast(Any, _NoQueryDB()),
        key=cast(Any, key),
        key_id="key-1",
        key_data=key_data,
    )

    assert prepared.update_data["name"] == "legacy-key"


def test_clear_oauth_invalid_response_invalidates_caches(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    cache_calls: list[tuple[str, str | None]] = []

    async def _fake_invalidate_key_cache(key_id: str) -> None:
        cache_calls.append(("key", key_id))

    async def _fake_invalidate_models_cache() -> None:
        cache_calls.append(("models", None))

    fake_provider_cache_module = types.ModuleType("src.services.cache.provider_cache")

    class _FakeProviderCacheService:
        @staticmethod
        async def invalidate_provider_api_key_cache(key_id: str) -> None:
            await _fake_invalidate_key_cache(key_id)

    setattr(fake_provider_cache_module, "ProviderCacheService", _FakeProviderCacheService)
    monkeypatch.setitem(
        sys.modules, "src.services.cache.provider_cache", fake_provider_cache_module
    )

    fake_models_service_module = types.ModuleType("src.services.cache.model_list_cache")
    setattr(
        fake_models_service_module, "invalidate_models_list_cache", _fake_invalidate_models_cache
    )
    monkeypatch.setitem(
        sys.modules, "src.services.cache.model_list_cache", fake_models_service_module
    )

    key = SimpleNamespace(
        oauth_invalid_at=datetime.now(timezone.utc),
        oauth_invalid_reason="forbidden",
        is_active=False,
    )
    db = _FakeClearOAuthDB(key=key)

    result = command_module.clear_oauth_invalid_response(cast(Any, db), key_id="key-1")

    assert result["message"] == "已清除 OAuth 失效标记"
    assert key.oauth_invalid_at is None
    assert key.oauth_invalid_reason is None
    assert key.is_active is False
    assert db.commit_count == 1
    assert cache_calls == [("key", "key-1"), ("models", None)]


def test_clear_oauth_invalid_response_noops_without_invalid_marker(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    cache_calls: list[tuple[str, str | None]] = []

    async def _fake_invalidate_key_cache(key_id: str) -> None:
        cache_calls.append(("key", key_id))

    async def _fake_invalidate_models_cache() -> None:
        cache_calls.append(("models", None))

    fake_provider_cache_module = types.ModuleType("src.services.cache.provider_cache")

    class _FakeProviderCacheService:
        @staticmethod
        async def invalidate_provider_api_key_cache(key_id: str) -> None:
            await _fake_invalidate_key_cache(key_id)

    setattr(fake_provider_cache_module, "ProviderCacheService", _FakeProviderCacheService)
    monkeypatch.setitem(
        sys.modules, "src.services.cache.provider_cache", fake_provider_cache_module
    )

    fake_models_service_module = types.ModuleType("src.services.cache.model_list_cache")
    setattr(
        fake_models_service_module, "invalidate_models_list_cache", _fake_invalidate_models_cache
    )
    monkeypatch.setitem(
        sys.modules, "src.services.cache.model_list_cache", fake_models_service_module
    )

    key = SimpleNamespace(
        oauth_invalid_at=None,
        oauth_invalid_reason=None,
        is_active=False,
    )
    db = _FakeClearOAuthDB(key=key)

    result = command_module.clear_oauth_invalid_response(cast(Any, db), key_id="key-1")

    assert result["message"] == "该 Key 当前无失效标记，无需清除"
    assert key.oauth_invalid_at is None
    assert key.oauth_invalid_reason is None
    assert key.is_active is False
    assert db.commit_count == 0
    assert cache_calls == []


@pytest.mark.asyncio
async def test_run_delete_key_side_effects_skip_disassociate(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    captured: dict[str, Any] = {}

    async def _fake_on_key_allowed_models_changed(**kwargs: Any) -> None:
        captured.update(kwargs)

    from src.services.model import global_model as global_model_module

    monkeypatch.setattr(
        global_model_module,
        "on_key_allowed_models_changed",
        _fake_on_key_allowed_models_changed,
    )

    await side_effects_module.run_delete_key_side_effects(
        db=cast(Any, object()),
        provider_id="provider-1",
        deleted_key_allowed_models=None,
    )

    assert captured["provider_id"] == "provider-1"
    assert captured["skip_disassociate"] is True


def test_cleanup_key_references_preserves_usage_and_video_tasks(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class _FakeDeleteStatement:
        def __init__(self, model: Any) -> None:
            self.kind = "delete"
            self.model = model
            self.conditions: list[Any] = []

        def where(self, *conditions: Any) -> _FakeDeleteStatement:
            self.conditions.extend(conditions)
            return self

    class _FakeUpdateStatement:
        def __init__(self, model: Any) -> None:
            self.kind = "update"
            self.model = model
            self.conditions: list[Any] = []
            self.values_dict: dict[str, Any] = {}

        def where(self, *conditions: Any) -> _FakeUpdateStatement:
            self.conditions.extend(conditions)
            return self

        def values(self, **values: Any) -> _FakeUpdateStatement:
            self.values_dict.update(values)
            return self

    class _FakeDB:
        def __init__(self) -> None:
            self.statements: list[Any] = []

        def execute(self, statement: Any) -> None:
            self.statements.append(statement)

        def get_bind(self) -> Any:
            return SimpleNamespace(dialect=SimpleNamespace(name="postgresql"))

    monkeypatch.setattr(side_effects_module, "sa_delete", lambda model: _FakeDeleteStatement(model))
    monkeypatch.setattr(side_effects_module, "sa_update", lambda model: _FakeUpdateStatement(model))

    db = _FakeDB()
    side_effects_module.cleanup_key_references(cast(Any, db), ["key-1", "key-2"])

    assert [(stmt.kind, stmt.model.__name__) for stmt in db.statements] == [
        ("delete", "GeminiFileMapping"),
        ("update", "Usage"),
        ("update", "VideoTask"),
    ]
    assert db.statements[1].values_dict == {"provider_api_key_id": None}
    assert db.statements[2].values_dict == {"key_id": None}


@pytest.mark.asyncio
async def test_batch_delete_endpoint_keys_response_cleans_related_references(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    cleanup_calls: list[list[str]] = []
    side_effect_calls: list[str | None] = []

    class _FakeBatchDeleteDB:
        def __init__(self, keys: list[Any]) -> None:
            self._keys = keys
            self.commit_count = 0
            self.executed: list[Any] = []

        def query(self, model: Any) -> _FakeQueryAll:
            _ = model
            return _FakeQueryAll(self._keys)

        def execute(self, statement: Any) -> None:
            self.executed.append(statement)

        def commit(self) -> None:
            self.commit_count += 1

        def rollback(self) -> None:
            raise AssertionError("rollback should not be called")

    async def _fake_run_delete_key_side_effects(
        db: Any,
        provider_id: str | None,
        deleted_key_allowed_models: list[str] | None,
    ) -> None:
        _ = db, deleted_key_allowed_models
        side_effect_calls.append(provider_id)

    monkeypatch.setattr(
        command_module,
        "cleanup_key_references",
        lambda _db, key_ids: cleanup_calls.append(list(key_ids)),
    )
    monkeypatch.setattr(
        command_module,
        "run_delete_key_side_effects",
        _fake_run_delete_key_side_effects,
    )

    keys = [
        SimpleNamespace(id="key-1", provider_id="provider-1"),
        SimpleNamespace(id="key-2", provider_id="provider-1"),
    ]
    db = _FakeBatchDeleteDB(keys)

    @contextmanager
    def _fake_get_db_context() -> Any:
        yield db

    monkeypatch.setattr(command_module, "get_db_context", _fake_get_db_context)

    result = await command_module.batch_delete_endpoint_keys_response(
        cast(Any, db),
        ["key-1", "key-2"],
    )

    assert result["success_count"] == 2
    assert result["failed_count"] == 0
    assert cleanup_calls == [["key-1", "key-2"]] or cleanup_calls == [["key-2", "key-1"]]
    assert side_effect_calls == ["provider-1"]
    assert db.commit_count == 1
    assert len(db.executed) == 1