gibby/tests/test_account_ops.py

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,
        )