from __future__ import annotations import time from typing import Any, cast import pytest from gibby.account_ops import ( PermanentAccountFailure, refresh_account_usage, snapshot_is_exhausted, window_used_percent, ) from gibby.client import OpenAIAPIError from gibby.models import ( UNKNOWN_EXHAUSTED_BACKOFF_SECONDS, AccountRecord, UsageSnapshot, UsageWindow, compute_cooldown_until, ) class FakeClient: def __init__(self, usage: UsageSnapshot, *, permanent: bool = False): self.usage = usage self.permanent = permanent self.refresh_calls: list[str] = [] async def refresh_access_token(self, refresh_token: str): self.refresh_calls.append(refresh_token) return ("new-token", "new-refresh", int(time.time()) + 600) async def fetch_usage_payload(self, access_token: str) -> dict: if self.permanent: raise OpenAIAPIError("invalid_grant", permanent=True, status_code=401) primary_window = self.usage.primary_window assert primary_window is not None secondary_window = ( { "used_percent": self.usage.secondary_window.used_percent, "limit_window_seconds": self.usage.secondary_window.limit_window_seconds, "reset_after_seconds": self.usage.secondary_window.reset_after_seconds, "reset_at": self.usage.secondary_window.reset_at, } if self.usage.secondary_window is not None else None ) return { "email": "acc@example.com", "account_id": "acc-1", "rate_limit": { "allowed": self.usage.allowed, "limit_reached": self.usage.limit_reached, "primary_window": { "used_percent": primary_window.used_percent, "limit_window_seconds": primary_window.limit_window_seconds, "reset_after_seconds": primary_window.reset_after_seconds, "reset_at": primary_window.reset_at, }, "secondary_window": secondary_window, }, } def make_usage( primary: int, secondary: int | None = None, *, limit_reached: bool = False, allowed: bool | None = None, reset_after: int = 10, checked_at: int | None = None, primary_reset_at: int | None = None, secondary_reset_at: int | None = None, primary_limit_window_seconds: int = 18000, secondary_limit_window_seconds: int = 604800, ) -> UsageSnapshot: checked_at = checked_at or int(time.time()) exhausted = ( primary >= 100 or (secondary is not None and secondary >= 100) or limit_reached ) if allowed is None: allowed = not exhausted return UsageSnapshot( checked_at=checked_at, used_percent=max(primary, secondary or 0), remaining_percent=max(0, 100 - max(primary, secondary or 0)), exhausted=exhausted or not allowed, primary_window=UsageWindow( primary, primary_limit_window_seconds, reset_after, primary_reset_at if primary_reset_at is not None else checked_at + reset_after, ), secondary_window=UsageWindow( secondary, secondary_limit_window_seconds, reset_after, secondary_reset_at if secondary_reset_at is not None else checked_at + reset_after, ) if secondary is not None else None, limit_reached=limit_reached, allowed=allowed, ) def test_window_used_percent_defaults_to_zero() -> None: assert window_used_percent(None) == 0 def test_snapshot_is_exhausted_checks_various_conditions() -> None: assert snapshot_is_exhausted(make_usage(94, 95), 96) is False assert snapshot_is_exhausted(make_usage(94), 95) is False assert snapshot_is_exhausted(make_usage(95), 95) is True assert snapshot_is_exhausted(make_usage(50, 95), 95) is True assert snapshot_is_exhausted(make_usage(50, limit_reached=True), 95) is True def test_compute_cooldown_until_uses_primary_blocking_window() -> None: usage = make_usage( 95, 20, checked_at=1000, primary_reset_at=1100, secondary_reset_at=1300 ) assert compute_cooldown_until(usage, 95) == 1100 def test_compute_cooldown_until_uses_secondary_blocking_window() -> None: usage = make_usage( 20, 95, checked_at=1000, primary_reset_at=1100, secondary_reset_at=1300 ) assert compute_cooldown_until(usage, 95) == 1300 def test_compute_cooldown_until_waits_for_all_blocking_windows() -> None: usage = make_usage( 95, 95, checked_at=1000, primary_reset_at=1100, secondary_reset_at=1300 ) assert compute_cooldown_until(usage, 95) == 1300 def test_compute_cooldown_until_uses_latest_window_when_blocker_is_ambiguous() -> None: usage = make_usage( 80, 40, limit_reached=True, allowed=False, checked_at=1000, primary_reset_at=1100, secondary_reset_at=1300, ) assert compute_cooldown_until(usage, 95) == 1300 def test_compute_cooldown_until_falls_back_to_limit_window_seconds() -> None: usage = make_usage( 95, checked_at=1000, reset_after=0, primary_reset_at=0, primary_limit_window_seconds=600, ) assert compute_cooldown_until(usage, 95) == 1600 def test_compute_cooldown_until_uses_backoff_when_no_reset_metadata_exists() -> None: usage = make_usage( 95, checked_at=1000, reset_after=0, primary_reset_at=0, primary_limit_window_seconds=0, ) assert compute_cooldown_until(usage, 95) == 1000 + UNKNOWN_EXHAUSTED_BACKOFF_SECONDS @pytest.mark.asyncio async def test_refresh_account_usage_populates_snapshot_and_cooldown() -> None: current = int(time.time()) account = AccountRecord( id="a1", access_token="tok", refresh_token="ref", expires_at=int(time.time()) + 600, ) usage = make_usage( 96, 1, limit_reached=True, allowed=False, checked_at=current, primary_reset_at=current + 100, secondary_reset_at=current + 300, ) client = FakeClient(usage) result = await refresh_account_usage(account, cast(Any, client), 95) assert result.used_percent == usage.used_percent assert result.limit_reached is usage.limit_reached assert result.allowed is usage.allowed assert account.last_known_usage is not None assert account.last_known_usage.used_percent == usage.used_percent assert account.cooldown_until == current + 100 assert account.email == "acc@example.com" assert account.id == "acc@example.com" @pytest.mark.asyncio async def test_refresh_account_usage_raises_permanent_failure() -> None: account = AccountRecord( id="a1", access_token="tok", refresh_token="ref", expires_at=int(time.time()) + 600, ) with pytest.raises(PermanentAccountFailure): await refresh_account_usage( account, cast(Any, FakeClient(make_usage(1), permanent=True)), 95, )