Aether/tests/services/test_quota_reader.py

from __future__ import annotations

from types import SimpleNamespace

import pytest

from src.services.provider.pool.account_state import resolve_pool_account_state
from src.services.provider.pool.dimensions._helpers import (
    extract_plan_type,
    extract_reset_seconds,
    extract_usage_ratio,
)
from src.services.provider_keys import quota_reader
from src.services.provider_keys.quota_reader import get_quota_reader
from src.services.scheduling.quota_skipper import is_key_quota_exhausted


def test_codex_reader_preserves_summary_formats() -> None:
    reader = get_quota_reader(
        "codex",
        {
            "codex": {
                "primary_used_percent": 14.8,
                "primary_reset_seconds": 266400,
                "secondary_used_percent": 27.9,
            }
        },
    )

    assert reader.display_summary() == "周剩余 85.2% (3天2小时后重置) | 5H剩余 72.1%"

    credits_reader = get_quota_reader(
        "codex",
        {"codex": {"has_credits": True, "credits_balance": 12.345}},
    )
    assert credits_reader.display_summary() == "积分 12.35"


def test_codex_reader_hides_reset_countdown_when_remaining_is_full() -> None:
    reader = get_quota_reader(
        "codex",
        {
            "codex": {
                "primary_used_percent": 0.0,
                "primary_reset_seconds": 266400,
                "secondary_used_percent": 0.0,
                "secondary_reset_seconds": 3600,
            }
        },
    )

    assert reader.display_summary() == "周剩余 100.0% | 5H剩余 100.0%"


def test_antigravity_reader_keeps_used_percent_fallbacks() -> None:
    reader = get_quota_reader(
        "antigravity",
        {
            "antigravity": {
                "quota_by_model": {
                    "claude-sonnet-4": {"used_percent": 100.0},
                    "gemini-2.5-pro": {"remaining_fraction": 0.6},
                }
            }
        },
    )

    exhausted = reader.is_exhausted("claude-sonnet-4")
    assert exhausted.exhausted is True
    assert exhausted.reason == "Antigravity 模型 claude-sonnet-4 配额剩余 0%"
    assert reader.display_summary() == "最低剩余 0.0% (2 模型)"
    assert reader.usage_ratio() == pytest.approx(0.7)


def test_dimension_helpers_delegate_to_unified_reader(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(quota_reader.time, "time", lambda: 100.0)
    key_obj = SimpleNamespace(
        upstream_metadata={
            "kiro": {
                "next_reset_at": 160.0,
                "usage_percentage": 45.0,
                "subscription_title": "Kiro Team",
            }
        },
        oauth_plan_type=None,
    )

    assert extract_plan_type(key_obj) == "team"
    assert extract_reset_seconds(key_obj) == pytest.approx(60.0)
    assert extract_usage_ratio(key_obj) == pytest.approx(0.45)


def test_extract_reset_seconds_uses_codex_weekly_reset_for_codex_provider() -> None:
    key_obj = SimpleNamespace(
        provider_type="codex",
        upstream_metadata={
            "codex": {
                "primary_reset_seconds": 1800.0,
                "secondary_reset_seconds": 120.0,
            }
        },
    )

    assert extract_reset_seconds(key_obj) == pytest.approx(1800.0)


def test_extract_reset_seconds_codex_weekly_full_quota_returns_none() -> None:
    key_obj = SimpleNamespace(
        provider_type="codex",
        upstream_metadata={
            "codex": {
                "primary_used_percent": 0.0,
                "primary_reset_seconds": 1800.0,
            }
        },
    )

    assert extract_reset_seconds(key_obj) is None


def test_extract_reset_seconds_prefers_codex_reset_at(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr("src.services.provider.pool.dimensions._helpers.time.time", lambda: 1000.0)
    key_obj = SimpleNamespace(
        provider_type="codex",
        upstream_metadata={
            "codex": {
                "primary_reset_at": 1200.0,
                "primary_reset_seconds": 9999.0,
            }
        },
    )

    assert extract_reset_seconds(key_obj) == pytest.approx(200.0)


def test_extract_reset_seconds_codex_fallback_corrects_elapsed(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr("src.services.provider.pool.dimensions._helpers.time.time", lambda: 2000.0)
    key_obj = SimpleNamespace(
        provider_type="codex",
        upstream_metadata={
            "codex": {
                "primary_reset_seconds": 180.0,
                "updated_at": 1900.0,
            }
        },
    )

    assert extract_reset_seconds(key_obj) == pytest.approx(80.0)


def test_extract_reset_seconds_codex_fallback_clamps_to_zero(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr("src.services.provider.pool.dimensions._helpers.time.time", lambda: 2000.0)
    key_obj = SimpleNamespace(
        provider_type="codex",
        upstream_metadata={
            "codex": {
                "primary_reset_seconds": 120.0,
                "updated_at": 1500.0,
            }
        },
    )

    assert extract_reset_seconds(key_obj) == pytest.approx(0.0)


def test_extract_reset_seconds_codex_fallback_future_updated_at_no_inflation(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr("src.services.provider.pool.dimensions._helpers.time.time", lambda: 2000.0)
    key_obj = SimpleNamespace(
        provider_type="codex",
        upstream_metadata={
            "codex": {
                "primary_reset_seconds": 120.0,
                "updated_at": 2600.0,
            }
        },
    )

    assert extract_reset_seconds(key_obj) == pytest.approx(120.0)


def test_resolve_pool_account_state_keeps_codex_metadata_block() -> None:
    state = resolve_pool_account_state(
        provider_type="codex",
        upstream_metadata={"codex": {"account_disabled": True, "message": "deactivated_workspace"}},
        oauth_invalid_reason=None,
    )

    assert state.blocked is True
    assert state.code == "workspace_deactivated"
    assert state.label == "工作区停用"
    assert state.reason == "deactivated_workspace"


def test_quota_skipper_uses_unified_reader() -> None:
    codex_key = SimpleNamespace(
        upstream_metadata={"codex": {"primary_used_percent": 100.0, "secondary_used_percent": 20.0}}
    )
    exhausted, reason = is_key_quota_exhausted("codex", codex_key, model_name="")  # type: ignore[arg-type]
    assert exhausted is True
    assert reason == "Codex 周限额剩余 0%"

    antigravity_key = SimpleNamespace(
        upstream_metadata={
            "antigravity": {"quota_by_model": {"gemini-2.5-pro": {"remaining_fraction": 0.0}}}
        }
    )
    exhausted, reason = is_key_quota_exhausted(
        "antigravity",
        antigravity_key,  # type: ignore[arg-type]
        model_name="gemini-2.5-pro",
    )
    assert exhausted is True
    assert reason == "Antigravity 模型 gemini-2.5-pro 配额剩余 0%"