Aether/tests/services/test_provider_keys_quota_cooldown.py

from __future__ import annotations

from types import SimpleNamespace
from typing import Any, cast

from src.core.provider_types import ProviderType
from src.services.provider_keys.quota_cooldown import resolve_effective_cooldown_reason


def _key_with_metadata(upstream_metadata: dict[str, Any]) -> Any:
    return cast(Any, SimpleNamespace(upstream_metadata=upstream_metadata))


def test_resolve_effective_cooldown_reason_prefers_redis_reason() -> None:
    key = _key_with_metadata({"codex": {"primary_used_percent": 100.0}})

    reason = resolve_effective_cooldown_reason(
        provider_type=ProviderType.CODEX,
        key=key,
        redis_reason="rate_limited_429",
    )

    assert reason == "rate_limited_429"


def test_resolve_effective_cooldown_reason_fallbacks_to_codex_quota_exhausted() -> None:
    key = _key_with_metadata({"codex": {"primary_used_percent": 100.0}})

    reason = resolve_effective_cooldown_reason(
        provider_type=ProviderType.CODEX,
        key=key,
        redis_reason=None,
    )

    assert reason == "quota_exhausted"


def test_resolve_effective_cooldown_reason_fallbacks_to_kiro_quota_exhausted() -> None:
    key = _key_with_metadata({"kiro": {"remaining": 0}})

    reason = resolve_effective_cooldown_reason(
        provider_type=ProviderType.KIRO,
        key=key,
        redis_reason=None,
    )

    assert reason == "quota_exhausted"


def test_resolve_effective_cooldown_reason_returns_none_when_not_exhausted() -> None:
    key = _key_with_metadata({"codex": {"primary_used_percent": 12.0}})

    reason = resolve_effective_cooldown_reason(
        provider_type=ProviderType.CODEX,
        key=key,
        redis_reason=None,
    )

    assert reason is None