Aether/tests/services/test_provider_keys_codex_realtime_quota.py

from __future__ import annotations

from types import SimpleNamespace
from typing import Any, cast

import pytest

from src.services.provider_keys import codex_realtime_quota as realtime_module
from src.services.provider_keys.codex_realtime_quota import (
    sync_codex_quota_from_response_headers,
)
from src.services.provider_keys.codex_usage_parser import (
    CodexUsageParseError,
    parse_codex_usage_headers,
)


class _FakeQuery:
    def __init__(self, db: "_FakeDB") -> None:
        self._db = db

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

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

    def first(self) -> Any:
        return self._db.key


class _FakeDB:
    def __init__(self, key: Any) -> None:
        self.key = key
        self.added: list[Any] = []
        self.query_count = 0

    def query(self, _model: Any) -> _FakeQuery:
        self.query_count += 1
        return _FakeQuery(self)

    def add(self, obj: Any) -> None:
        self.added.append(obj)


def _paid_headers(**overrides: Any) -> dict[str, Any]:
    base = {
        "x-codex-plan-type": "team",
        "x-codex-primary-used-percent": "3",
        "x-codex-secondary-used-percent": "64",
        "x-codex-primary-window-minutes": "300",
        "x-codex-secondary-window-minutes": "10080",
        "x-codex-primary-reset-after-seconds": "411",
        "x-codex-secondary-reset-after-seconds": "267545",
        "x-codex-primary-reset-at": "1772259405",
        "x-codex-secondary-reset-at": "1772526539",
        "x-codex-credits-has-credits": "False",
        "x-codex-credits-balance": "",
        "x-codex-credits-unlimited": "False",
    }
    base.update(overrides)
    return base


def _as_session(db: _FakeDB) -> Any:
    # 测试中使用最小假对象模拟 Session 接口
    return cast(Any, db)


def test_parse_codex_usage_headers_paid_windows_mapping() -> None:
    parsed = parse_codex_usage_headers(_paid_headers())
    assert parsed is not None
    assert parsed["plan_type"] == "team"
    # 与 wham 解析保持一致：primary=周限额，secondary=5H 限额
    assert parsed["primary_used_percent"] == 64.0
    assert parsed["secondary_used_percent"] == 3.0
    assert parsed["primary_window_minutes"] == 10080
    assert parsed["secondary_window_minutes"] == 300
    assert parsed["has_credits"] is False
    assert parsed["credits_unlimited"] is False
    assert "credits_balance" not in parsed


@pytest.mark.parametrize("plan_type", ["plus", "enterprise"])
def test_parse_codex_usage_headers_paid_plan_windows_mapping(plan_type: str) -> None:
    parsed = parse_codex_usage_headers(_paid_headers(**{"x-codex-plan-type": plan_type}))
    assert parsed is not None
    assert parsed["plan_type"] == plan_type
    assert parsed["primary_used_percent"] == 64.0
    assert parsed["secondary_used_percent"] == 3.0
    assert parsed["primary_window_minutes"] == 10080
    assert parsed["secondary_window_minutes"] == 300


def test_parse_codex_usage_headers_free_uses_primary_only() -> None:
    parsed = parse_codex_usage_headers(
        {
            "x-codex-plan-type": "FREE",
            "x-codex-primary-used-percent": "12.5",
            "x-codex-primary-window-minutes": "10080",
            "x-codex-primary-reset-after-seconds": "120",
            "x-codex-primary-reset-at": "1700000000",
        }
    )
    assert parsed is not None
    assert parsed["plan_type"] == "free"
    assert parsed["primary_used_percent"] == 12.5
    assert parsed["primary_window_minutes"] == 10080
    assert "secondary_used_percent" not in parsed


def test_parse_codex_usage_headers_invalid_type_raises() -> None:
    with pytest.raises(CodexUsageParseError, match="headers.primary_window.limit_window_minutes"):
        parse_codex_usage_headers({"x-codex-primary-window-minutes": "abc"})


def test_sync_codex_quota_from_headers_updates_and_preserves_existing_fields() -> None:
    realtime_module._header_fingerprint_cache.clear()

    key = SimpleNamespace(
        id="sync-update-key",
        provider=SimpleNamespace(provider_type="codex"),
        upstream_metadata={
            "codex": {
                "primary_used_percent": 50.0,
                "legacy_marker": "keep-me",
            }
        },
    )
    db = _FakeDB(key)

    updated = sync_codex_quota_from_response_headers(
        db=_as_session(db),
        provider_api_key_id="sync-update-key",
        response_headers=_paid_headers(),
    )

    assert updated is True
    assert db.added == [key]
    codex_meta = key.upstream_metadata["codex"]
    assert codex_meta["primary_used_percent"] == 64.0
    assert codex_meta["secondary_used_percent"] == 3.0
    # 旧字段应被保留（解析器只覆盖已知配额字段）
    assert codex_meta["legacy_marker"] == "keep-me"


def test_sync_codex_quota_from_headers_skips_when_only_reset_seconds_changed() -> None:
    realtime_module._header_fingerprint_cache.clear()

    key = SimpleNamespace(
        id="sync-volatile-key",
        provider=SimpleNamespace(provider_type="codex"),
        upstream_metadata={
            "codex": {
                "plan_type": "team",
                "primary_used_percent": 64.0,
                "secondary_used_percent": 3.0,
                "primary_window_minutes": 10080,
                "secondary_window_minutes": 300,
                "primary_reset_at": 1772526539,
                "secondary_reset_at": 1772259405,
                "primary_reset_seconds": 999999,
                "secondary_reset_seconds": 999999,
                "has_credits": False,
                "credits_unlimited": False,
            }
        },
    )
    db = _FakeDB(key)

    updated = sync_codex_quota_from_response_headers(
        db=_as_session(db),
        provider_api_key_id="sync-volatile-key",
        response_headers=_paid_headers(),
    )

    assert updated is False
    assert db.added == []


def test_sync_codex_quota_from_headers_cache_hit_skips_second_query() -> None:
    realtime_module._header_fingerprint_cache.clear()

    key = SimpleNamespace(
        id="sync-cache-key",
        provider=SimpleNamespace(provider_type="codex"),
        upstream_metadata={},
    )
    db = _FakeDB(key)
    headers = _paid_headers()

    first_updated = sync_codex_quota_from_response_headers(
        db=_as_session(db),
        provider_api_key_id="sync-cache-key",
        response_headers=headers,
    )
    second_updated = sync_codex_quota_from_response_headers(
        db=_as_session(db),
        provider_api_key_id="sync-cache-key",
        response_headers=headers,
    )

    assert first_updated is True
    assert second_updated is False
    assert db.query_count == 1


def test_sync_codex_quota_from_headers_non_codex_key_is_ignored() -> None:
    realtime_module._header_fingerprint_cache.clear()

    key = SimpleNamespace(
        id="sync-kiro-key",
        provider=SimpleNamespace(provider_type="kiro"),
        upstream_metadata={},
    )
    db = _FakeDB(key)

    updated = sync_codex_quota_from_response_headers(
        db=_as_session(db),
        provider_api_key_id="sync-kiro-key",
        response_headers=_paid_headers(),
    )

    assert updated is False
    assert db.added == []


def test_sync_codex_quota_from_headers_parse_error_is_non_blocking() -> None:
    realtime_module._header_fingerprint_cache.clear()

    key = SimpleNamespace(
        id="sync-parse-error-key",
        provider=SimpleNamespace(provider_type="codex"),
        upstream_metadata={},
    )
    db = _FakeDB(key)

    updated = sync_codex_quota_from_response_headers(
        db=_as_session(db),
        provider_api_key_id="sync-parse-error-key",
        response_headers={
            "x-codex-primary-window-minutes": "abc",
        },
    )

    assert updated is False
    assert db.query_count == 0
    assert db.added == []


def test_header_fingerprint_cache_prunes_expired_entries() -> None:
    realtime_module._header_fingerprint_cache.clear()

    realtime_module._set_cached_fingerprint("expired-key", "fp-1", now_ts=10.0)
    realtime_module._set_cached_fingerprint("fresh-key", "fp-2", now_ts=10.0 + 31.0)

    assert "expired-key" not in realtime_module._header_fingerprint_cache
    assert "fresh-key" in realtime_module._header_fingerprint_cache


def test_header_fingerprint_cache_respects_max_entries() -> None:
    realtime_module._header_fingerprint_cache.clear()

    max_entries = realtime_module._CACHE_MAX_ENTRIES
    for i in range(max_entries + 8):
        realtime_module._set_cached_fingerprint(f"key-{i}", f"fp-{i}", now_ts=1000.0)

    assert len(realtime_module._header_fingerprint_cache) == max_entries