1
0
Fork 0
This commit is contained in:
Arthur K. 2026-04-20 23:41:37 +03:00
commit 7cef56de15
23 changed files with 3136 additions and 0 deletions

233
tests/test_account_ops.py Normal file
View file

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

89
tests/test_app.py Normal file
View file

@ -0,0 +1,89 @@
from __future__ import annotations
from typing import Any
from fastapi.testclient import TestClient
from gibby.app import create_app
from gibby.manager import AccountManager, NoUsableAccountError
class StubManager(AccountManager):
def __init__(
self,
response: dict[str, Any] | None = None,
usage_response: dict[str, Any] | None = None,
error: Exception | None = None,
):
self.response = response
self.usage_response = usage_response
self.error = error
async def issue_token_response(self):
if self.error is not None:
raise self.error
if self.response is not None:
return self.response
return {}
async def get_usage_report(self):
if self.usage_response is not None:
return self.usage_response
return {"accounts": [], "active_account_id": None, "count": 0}
def test_health_ok() -> None:
client = TestClient(create_app(StubManager(response={})))
response = client.get("/health")
assert response.status_code == 200
assert response.text == "ok"
def test_token_success_shape() -> None:
payload = {
"token": "abc",
"limit": {
"used_percent": 10,
"remaining_percent": 90,
"exhausted": False,
"needs_prepare": False,
},
"usage": {"primary_window": None, "secondary_window": None},
}
client = TestClient(create_app(StubManager(response=payload)))
response = client.get("/token")
assert response.status_code == 200
assert response.json() == payload
def test_token_error_shape() -> None:
client = TestClient(
create_app(
StubManager(error=NoUsableAccountError("No usable account available"))
)
)
response = client.get("/token")
assert response.status_code == 503
assert response.json() == {"error": "No usable account available"}
def test_usage_success_shape() -> None:
payload = {
"accounts": [
{
"id": "a1",
"email": "a1@example.com",
"status": "ok",
"used_percent": 12,
"remaining_percent": 88,
"cooldown_until": None,
"primary_window": {"used_percent": 12},
"secondary_window": {"used_percent": 1},
}
],
"active_account_id": "a1",
"count": 1,
}
client = TestClient(create_app(StubManager(usage_response=payload)))
response = client.get("/usage")
assert response.status_code == 200
assert response.json() == payload

570
tests/test_core.py Normal file
View file

@ -0,0 +1,570 @@
from __future__ import annotations
import time
from pathlib import Path
import pytest
from gibby.client import OpenAIAPIError, OpenAIClient
from gibby.manager import AccountManager, NoUsableAccountError
from gibby.models import AccountRecord, StateFile, UsageSnapshot, UsageWindow
from gibby.settings import Settings
from gibby.store import JsonStateStore
class FakeClient(OpenAIClient):
def __init__(
self,
usage_by_token=None,
refresh_map=None,
failing_tokens=None,
permanent_refresh_tokens=None,
):
self.usage_by_token = usage_by_token or {}
self.refresh_map = refresh_map or {}
self.failing_tokens = set(failing_tokens or [])
self.permanent_refresh_tokens = set(permanent_refresh_tokens or [])
self.fetched_tokens: list[str] = []
self.refresh_calls: list[str] = []
async def refresh_access_token(self, refresh_token: str):
self.refresh_calls.append(refresh_token)
if refresh_token in self.permanent_refresh_tokens:
raise OpenAIAPIError("invalid_grant", permanent=True, status_code=401)
return self.refresh_map[refresh_token]
async def fetch_usage_payload(self, access_token: str):
self.fetched_tokens.append(access_token)
if access_token in self.failing_tokens:
raise RuntimeError("usage failed")
usage = self.usage_by_token[access_token]
return {
"email": f"{access_token}@example.com",
"account_id": f"acct-{access_token}",
"rate_limit": {
"allowed": usage.allowed,
"limit_reached": usage.limit_reached,
"primary_window": {
"used_percent": usage.primary_window.used_percent,
"limit_window_seconds": usage.primary_window.limit_window_seconds,
"reset_after_seconds": usage.primary_window.reset_after_seconds,
"reset_at": usage.primary_window.reset_at,
}
if usage.primary_window
else None,
"secondary_window": {
"used_percent": usage.secondary_window.used_percent,
"limit_window_seconds": usage.secondary_window.limit_window_seconds,
"reset_after_seconds": usage.secondary_window.reset_after_seconds,
"reset_at": usage.secondary_window.reset_at,
}
if usage.secondary_window
else None,
},
}
def make_usage(
*,
used: int,
secondary_used: int | None = None,
limit_reached: bool = False,
reset_after: int = 0,
) -> UsageSnapshot:
exhausted = (
used >= 100
or (secondary_used is not None and secondary_used >= 100)
or limit_reached
)
return UsageSnapshot(
checked_at=int(time.time()),
used_percent=used,
remaining_percent=max(0, 100 - used),
exhausted=exhausted,
primary_window=UsageWindow(
used_percent=used,
limit_window_seconds=604800,
reset_after_seconds=reset_after,
reset_at=int(time.time()) + reset_after if reset_after else 0,
),
secondary_window=UsageWindow(
used_percent=secondary_used,
limit_window_seconds=604800,
reset_after_seconds=reset_after,
reset_at=int(time.time()) + reset_after if reset_after else 0,
)
if secondary_used is not None
else None,
limit_reached=limit_reached,
allowed=not exhausted,
)
def make_manager(
store: JsonStateStore,
client: FakeClient,
*,
threshold: int = 95,
) -> AccountManager:
return AccountManager(
store,
client,
Settings(data_dir=store.path.parent, exhausted_usage_threshold=threshold),
)
def make_store(tmp_path: Path, state: StateFile) -> JsonStateStore:
store = JsonStateStore(tmp_path / "accounts.json")
store.save(state)
return store
@pytest.mark.asyncio
async def test_prefers_active_account_when_locally_usable(tmp_path: Path) -> None:
active = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=20),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=70),
)
store = make_store(
tmp_path, StateFile(active_account_id="a1", accounts=[active, second])
)
client = FakeClient(
usage_by_token={
"tok-a1": make_usage(used=21),
"tok-a2": make_usage(used=72),
}
)
payload = await make_manager(store, client).issue_token_response()
assert payload["token"] == "tok-a1"
assert client.fetched_tokens == ["tok-a1"]
assert store.load().active_account_id == "tok-a1@example.com"
@pytest.mark.asyncio
async def test_prefers_higher_primary_usage_from_saved_snapshot(tmp_path: Path) -> None:
first = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=20),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=70),
)
store = make_store(tmp_path, StateFile(accounts=[first, second]))
client = FakeClient(usage_by_token={"tok-a2": make_usage(used=72)})
payload = await make_manager(store, client).issue_token_response()
assert payload["token"] == "tok-a2"
assert client.fetched_tokens == ["tok-a2"]
saved = store.load()
assert saved.active_account_id == "tok-a2@example.com"
@pytest.mark.asyncio
async def test_breaks_ties_with_secondary_usage(tmp_path: Path) -> None:
first = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=60, secondary_used=10),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=60, secondary_used=40),
)
store = make_store(tmp_path, StateFile(accounts=[first, second]))
client = FakeClient(
usage_by_token={"tok-a2": make_usage(used=61, secondary_used=41)}
)
payload = await make_manager(store, client).issue_token_response()
assert payload["token"] == "tok-a2"
assert client.fetched_tokens == ["tok-a2"]
@pytest.mark.asyncio
async def test_treats_missing_secondary_as_zero(tmp_path: Path) -> None:
first = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=60),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=60, secondary_used=1),
)
store = make_store(tmp_path, StateFile(accounts=[first, second]))
client = FakeClient(
usage_by_token={"tok-a2": make_usage(used=61, secondary_used=1)}
)
payload = await make_manager(store, client).issue_token_response()
assert payload["token"] == "tok-a2"
assert client.fetched_tokens == ["tok-a2"]
@pytest.mark.asyncio
async def test_skips_account_still_in_cooldown(tmp_path: Path) -> None:
active = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
cooldown_until=int(time.time()) + 300,
last_known_usage=make_usage(used=80),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=20),
)
store = make_store(
tmp_path, StateFile(active_account_id="a1", accounts=[active, second])
)
client = FakeClient(usage_by_token={"tok-a2": make_usage(used=25)})
payload = await make_manager(store, client).issue_token_response()
assert payload["token"] == "tok-a2"
assert client.fetched_tokens == ["tok-a2"]
assert store.load().active_account_id == "tok-a2@example.com"
@pytest.mark.asyncio
async def test_skips_active_account_blocked_by_local_exhausted_snapshot(
tmp_path: Path,
) -> None:
active = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=96),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=20),
)
store = make_store(
tmp_path, StateFile(active_account_id="a1", accounts=[active, second])
)
client = FakeClient(usage_by_token={"tok-a2": make_usage(used=25)})
payload = await make_manager(store, client).issue_token_response()
assert payload["token"] == "tok-a2"
assert client.fetched_tokens == ["tok-a2"]
assert store.load().active_account_id == "tok-a2@example.com"
@pytest.mark.asyncio
async def test_live_checks_depleted_and_moves_to_next(tmp_path: Path) -> None:
high = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=94),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=50),
)
store = make_store(tmp_path, StateFile(accounts=[high, second]))
client = FakeClient(
usage_by_token={
"tok-a1": make_usage(used=96, reset_after=120),
"tok-a2": make_usage(used=52),
}
)
payload = await make_manager(store, client).issue_token_response()
assert payload["token"] == "tok-a2"
assert client.fetched_tokens == ["tok-a1", "tok-a2"]
@pytest.mark.asyncio
async def test_live_checks_secondary_depleted_and_moves_to_next(tmp_path: Path) -> None:
high = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=30, secondary_used=94),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=20, secondary_used=10),
)
store = make_store(tmp_path, StateFile(accounts=[high, second]))
client = FakeClient(
usage_by_token={
"tok-a1": make_usage(used=30, secondary_used=100, reset_after=120),
"tok-a2": make_usage(used=22, secondary_used=10),
}
)
payload = await make_manager(store, client).issue_token_response()
assert payload["token"] == "tok-a2"
assert client.fetched_tokens == ["tok-a1", "tok-a2"]
@pytest.mark.asyncio
async def test_falls_through_when_live_usage_is_depleted(tmp_path: Path) -> None:
first = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=80),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=70),
)
store = make_store(tmp_path, StateFile(accounts=[first, second]))
client = FakeClient(
usage_by_token={
"tok-a1": make_usage(used=95, reset_after=120),
"tok-a2": make_usage(used=71),
}
)
payload = await make_manager(store, client).issue_token_response()
saved = store.load()
assert payload["token"] == "tok-a2"
assert client.fetched_tokens == ["tok-a1", "tok-a2"]
depleted = next(
account for account in saved.accounts if account.id == "tok-a1@example.com"
)
assert depleted.cooldown_until is not None
assert depleted.last_known_usage is not None
assert depleted.last_known_usage.used_percent == 95
@pytest.mark.asyncio
async def test_keeps_account_out_until_blocking_window_resets(tmp_path: Path) -> None:
current = int(time.time())
blocked = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=80, secondary_used=40),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=70),
)
store = make_store(tmp_path, StateFile(accounts=[blocked, second]))
blocked_usage = UsageSnapshot(
checked_at=current,
used_percent=80,
remaining_percent=20,
exhausted=True,
primary_window=UsageWindow(
used_percent=80,
limit_window_seconds=604800,
reset_after_seconds=60,
reset_at=current + 60,
),
secondary_window=UsageWindow(
used_percent=40,
limit_window_seconds=604800,
reset_after_seconds=240,
reset_at=current + 240,
),
limit_reached=True,
allowed=False,
)
client = FakeClient(
usage_by_token={
"tok-a1": blocked_usage,
"tok-a2": make_usage(used=71),
}
)
payload = await make_manager(store, client).issue_token_response()
saved = store.load()
blocked_saved = next(
account for account in saved.accounts if account.id == "tok-a1@example.com"
)
assert payload["token"] == "tok-a2"
assert blocked_saved.cooldown_until == current + 240
assert client.fetched_tokens == ["tok-a1", "tok-a2"]
@pytest.mark.asyncio
async def test_refreshes_expired_token_before_usage(tmp_path: Path) -> None:
active = AccountRecord(
id="a1",
access_token="old-token",
refresh_token="ref-a1",
expires_at=int(time.time()) - 1,
last_known_usage=make_usage(used=20),
)
store = make_store(tmp_path, StateFile(active_account_id="a1", accounts=[active]))
client = FakeClient(
usage_by_token={"new-token": make_usage(used=20)},
refresh_map={"ref-a1": ("new-token", "new-refresh", int(time.time()) + 600)},
)
payload = await make_manager(store, client).issue_token_response()
saved = store.load()
assert payload["token"] == "new-token"
assert client.refresh_calls == ["ref-a1"]
assert client.fetched_tokens == ["new-token"]
assert saved.accounts[0].id == "new-token@example.com"
assert saved.accounts[0].access_token == "new-token"
assert saved.accounts[0].refresh_token == "new-refresh"
@pytest.mark.asyncio
async def test_raises_when_all_accounts_unusable(tmp_path: Path) -> None:
active = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=80),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=70),
)
store = make_store(
tmp_path, StateFile(active_account_id="a1", accounts=[active, second])
)
client = FakeClient(
usage_by_token={
"tok-a1": make_usage(used=96, reset_after=120),
"tok-a2": make_usage(used=97, reset_after=120),
}
)
with pytest.raises(NoUsableAccountError):
await make_manager(store, client).issue_token_response()
saved = store.load()
assert all(account.cooldown_until is not None for account in saved.accounts)
@pytest.mark.asyncio
async def test_threshold_can_be_overridden_for_selection(tmp_path: Path) -> None:
active = AccountRecord(
id="a1",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=96),
)
second = AccountRecord(
id="a2",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=20),
)
store = make_store(
tmp_path, StateFile(active_account_id="a1", accounts=[active, second])
)
client = FakeClient(
usage_by_token={
"tok-a1": make_usage(used=96),
"tok-a2": make_usage(used=25),
}
)
payload = await make_manager(store, client, threshold=97).issue_token_response()
assert payload["token"] == "tok-a1"
assert client.fetched_tokens == ["tok-a1"]
@pytest.mark.asyncio
async def test_removes_account_and_records_failed_email_on_permanent_refresh_failure(
tmp_path: Path,
) -> None:
dead = AccountRecord(
id="a1",
email="dead@example.com",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) - 1,
last_known_usage=make_usage(used=80),
)
alive = AccountRecord(
id="a2",
email="alive@example.com",
access_token="tok-a2",
refresh_token="ref-a2",
expires_at=int(time.time()) + 600,
last_known_usage=make_usage(used=70),
)
store = make_store(tmp_path, StateFile(accounts=[dead, alive]))
client = FakeClient(
usage_by_token={"tok-a2": make_usage(used=71)},
permanent_refresh_tokens={"ref-a1"},
)
payload = await make_manager(store, client).issue_token_response()
assert payload["token"] == "tok-a2"
saved = store.load()
assert [account.id for account in saved.accounts] == ["tok-a2@example.com"]
assert (tmp_path / "failed.txt").read_text().splitlines() == ["dead@example.com"]

177
tests/test_oauth_helper.py Normal file
View file

@ -0,0 +1,177 @@
from __future__ import annotations
import asyncio
import sys
import urllib.parse
from pathlib import Path
from typing import Any, cast
import httpx
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts"))
from gibby.client import OpenAIAPIError
from gibby.models import UsageSnapshot, UsageWindow
from gibby.oauth import build_authorize_url, generate_pkce_pair
from gibby.settings import Settings
from oauth_helper import ( # type: ignore[import-not-found]
exchange_and_store_account,
parse_redirect_url,
wait_for_callback,
)
from gibby.store import JsonStateStore
class FakeClient:
def __init__(
self,
settings: Settings,
usage: UsageSnapshot,
*,
transient_usage_failure: bool = False,
):
self.settings = settings
self.usage = usage
self.transient_usage_failure = transient_usage_failure
async def exchange_code(self, code: str, verifier: str) -> tuple[str, str, int]:
return ("access-token", "refresh-token", 1776000000)
async def refresh_access_token(self, refresh_token: str) -> tuple[str, str, int]:
return ("access-token", refresh_token, 1776000000)
async def fetch_usage_payload(self, access_token: str) -> dict:
if self.transient_usage_failure:
raise OpenAIAPIError("usage timeout", permanent=False)
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": "oauth@example.com",
"account_id": "oauth-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) -> UsageSnapshot:
return UsageSnapshot(
checked_at=1775000000,
used_percent=max(primary, secondary or 0),
remaining_percent=max(0, 100 - max(primary, secondary or 0)),
exhausted=False,
primary_window=UsageWindow(primary, 18000, 100, 1775000100),
secondary_window=UsageWindow(secondary, 604800, 100, 1775000100)
if secondary is not None
else None,
limit_reached=False,
allowed=True,
)
def test_generate_pkce_pair_shapes() -> None:
verifier, challenge = generate_pkce_pair()
assert len(verifier) > 40
assert len(challenge) > 40
assert "=" not in challenge
def test_build_authorize_url_contains_redirect_and_state() -> None:
settings = Settings(callback_host="localhost", callback_port=1455)
url = build_authorize_url(settings, "challenge", "state-123")
parsed = urllib.parse.urlparse(url)
query = urllib.parse.parse_qs(parsed.query)
assert parsed.scheme == "https"
assert query["redirect_uri"] == ["http://localhost:1455/auth/callback"]
assert query["state"] == ["state-123"]
assert query["code_challenge"] == ["challenge"]
def test_parse_redirect_url_extracts_code_and_state() -> None:
code, state = parse_redirect_url(
"http://127.0.0.1:1455/auth/callback?code=abc&state=xyz"
)
assert code == "abc"
assert state == "xyz"
@pytest.mark.asyncio
async def test_wait_for_callback_receives_code_and_state() -> None:
task = asyncio.create_task(
wait_for_callback("127.0.0.1", 18555, "state-1", timeout=5)
)
await asyncio.sleep(0.05)
async with httpx.AsyncClient() as client:
response = await client.get(
"http://127.0.0.1:18555/auth/callback?code=abc123&state=state-1",
timeout=5,
)
result = await task
assert response.status_code == 200
assert result == ("abc123", "state-1")
@pytest.mark.asyncio
async def test_exchange_and_store_account_populates_usage_snapshot(
tmp_path: Path,
) -> None:
store = JsonStateStore(tmp_path / "accounts.json")
settings = Settings(data_dir=tmp_path)
client = FakeClient(settings, make_usage(12, 3))
account = await exchange_and_store_account(
store,
cast(Any, client),
"code",
"verifier",
False,
)
assert account.last_known_usage is not None
assert account.id == "oauth@example.com"
assert account.last_known_usage.primary_window is not None
assert account.last_known_usage.primary_window.used_percent == 12
assert account.last_known_usage.secondary_window is not None
assert account.last_known_usage.secondary_window.used_percent == 3
@pytest.mark.asyncio
async def test_exchange_and_store_account_keeps_tokens_on_transient_usage_failure(
tmp_path: Path,
) -> None:
store = JsonStateStore(tmp_path / "accounts.json")
settings = Settings(data_dir=tmp_path)
client = FakeClient(settings, make_usage(12, 3), transient_usage_failure=True)
account = await exchange_and_store_account(
store,
cast(Any, client),
"code",
"verifier",
False,
)
saved = store.load()
assert account.last_known_usage is None
assert saved.accounts[0].access_token == "access-token"
assert saved.accounts[0].last_error is not None

View file

@ -0,0 +1,132 @@
from __future__ import annotations
import time
import sys
from pathlib import Path
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts"))
from gibby.client import OpenAIAPIError
from gibby.models import AccountRecord, StateFile, UsageSnapshot, UsageWindow
from gibby.store import JsonStateStore
import refresh_limits # type: ignore[import-not-found]
def make_usage(primary: int, secondary: int | None = None) -> UsageSnapshot:
return UsageSnapshot(
checked_at=int(time.time()),
used_percent=max(primary, secondary or 0),
remaining_percent=max(0, 100 - max(primary, secondary or 0)),
exhausted=False,
primary_window=UsageWindow(primary, 18000, 10, int(time.time()) + 10),
secondary_window=UsageWindow(secondary, 604800, 10, int(time.time()) + 10)
if secondary is not None
else None,
limit_reached=False,
allowed=True,
)
class FakeClient:
def __init__(self, settings, *, permanent: bool = False):
self.settings = settings
self.permanent = permanent
async def aclose(self) -> None:
return
async def refresh_access_token(self, refresh_token: str):
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)
usage = make_usage(12, 4)
primary_window = usage.primary_window
assert primary_window is not None
secondary_window = (
{
"used_percent": usage.secondary_window.used_percent,
"limit_window_seconds": usage.secondary_window.limit_window_seconds,
"reset_after_seconds": usage.secondary_window.reset_after_seconds,
"reset_at": usage.secondary_window.reset_at,
}
if usage.secondary_window is not None
else None
)
return {
"email": "acc@example.com",
"account_id": "acc-1",
"rate_limit": {
"allowed": usage.allowed,
"limit_reached": 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,
},
}
@pytest.mark.asyncio
async def test_refresh_limits_updates_all_accounts(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
store = JsonStateStore(tmp_path / "accounts.json")
store.save(
StateFile(
accounts=[
AccountRecord(
id="acc@example.com",
email="acc@example.com",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
)
]
)
)
monkeypatch.setattr(refresh_limits, "OpenAIClient", FakeClient)
await refresh_limits.run(tmp_path)
state = store.load()
assert state.accounts[0].last_known_usage is not None
assert state.accounts[0].last_known_usage.primary_window is not None
assert state.accounts[0].last_known_usage.primary_window.used_percent == 12
@pytest.mark.asyncio
async def test_refresh_limits_removes_permanently_failed_account(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
store = JsonStateStore(tmp_path / "accounts.json")
store.save(
StateFile(
accounts=[
AccountRecord(
id="dead@example.com",
email="dead@example.com",
access_token="tok-a1",
refresh_token="ref-a1",
expires_at=int(time.time()) + 600,
)
]
)
)
def permanent_client(settings):
return FakeClient(settings, permanent=True)
monkeypatch.setattr(refresh_limits, "OpenAIClient", permanent_client)
await refresh_limits.run(tmp_path)
state = store.load()
assert state.accounts == []
assert (tmp_path / "failed.txt").read_text().splitlines() == ["dead@example.com"]

96
tests/test_store.py Normal file
View file

@ -0,0 +1,96 @@
from __future__ import annotations
import json
from gibby.models import AccountRecord, StateFile, UsageSnapshot, UsageWindow
from gibby.store import JsonStateStore
def test_store_writes_canonical_usage_snapshot_shape(tmp_path) -> None:
store = JsonStateStore(tmp_path / "accounts.json")
snapshot = UsageSnapshot(
checked_at=1000,
used_percent=75,
remaining_percent=25,
exhausted=False,
primary_window=UsageWindow(75, 18000, 300, 1300),
secondary_window=UsageWindow(10, 604800, 3600, 4600),
limit_reached=False,
allowed=True,
)
store.save(
StateFile(
accounts=[
AccountRecord(
id="acc@example.com",
email="acc@example.com",
access_token="tok",
refresh_token="ref",
expires_at=2000,
last_known_usage=snapshot,
)
]
)
)
payload = json.loads((tmp_path / "accounts.json").read_text())
saved_snapshot = payload["accounts"][0]["last_known_usage"]
assert set(saved_snapshot) == {
"checked_at",
"primary_window",
"secondary_window",
"limit_reached",
"allowed",
}
def test_store_load_reconstructs_derived_usage_fields(tmp_path) -> None:
path = tmp_path / "accounts.json"
path.write_text(
json.dumps(
{
"version": 1,
"active_account_id": "acc@example.com",
"accounts": [
{
"id": "acc@example.com",
"email": "acc@example.com",
"account_id": "acc-1",
"access_token": "tok",
"refresh_token": "ref",
"expires_at": 2000,
"cooldown_until": None,
"last_known_usage": {
"checked_at": 1000,
"primary_window": {
"used_percent": 80,
"limit_window_seconds": 18000,
"reset_after_seconds": 300,
"reset_at": 1300,
},
"secondary_window": {
"used_percent": 100,
"limit_window_seconds": 604800,
"reset_after_seconds": 3600,
"reset_at": 4600,
},
"limit_reached": False,
"allowed": True,
},
"last_error": None,
}
],
}
)
)
state = JsonStateStore(path).load()
snapshot = state.accounts[0].last_known_usage
assert snapshot is not None
assert snapshot.used_percent == 100
assert snapshot.remaining_percent == 0
assert snapshot.exhausted is True
assert snapshot.limit_reached is False
assert snapshot.allowed is True