tests/test_account_ops.py

@@ -0,0 +1,233 @@
+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,
+        )