This commit is contained in:
commit
7cef56de15
23 changed files with 3136 additions and 0 deletions
233
tests/test_account_ops.py
Normal file
233
tests/test_account_ops.py
Normal 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
89
tests/test_app.py
Normal 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
570
tests/test_core.py
Normal 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
177
tests/test_oauth_helper.py
Normal 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
|
||||
132
tests/test_refresh_limits.py
Normal file
132
tests/test_refresh_limits.py
Normal 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
96
tests/test_store.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue