266 lines
9.5 KiB
Python
266 lines
9.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
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,
|
|
invalid_tokens=None,
|
|
permanent_refresh_tokens=None,
|
|
):
|
|
self.usage_by_token = usage_by_token or {}
|
|
self.refresh_map = refresh_map or {}
|
|
self.invalid_tokens = set(invalid_tokens or [])
|
|
self.permanent_refresh_tokens = set(permanent_refresh_tokens or [])
|
|
self.fetched_usage_tokens: list[str] = []
|
|
self.validated_tokens: list[str] = []
|
|
self.refresh_calls: list[str] = []
|
|
self.settings = Settings(data_dir=Path("."))
|
|
|
|
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_usage_tokens.append(access_token)
|
|
usage = self.usage_by_token[access_token]
|
|
return {
|
|
"email": f"{access_token}@example.com",
|
|
"rate_limit": {
|
|
"allowed": usage.allowed,
|
|
"limit_reached": usage.limit_reached,
|
|
"primary_window": {
|
|
"used_percent": usage.primary_window.used_percent,
|
|
"reset_at": usage.primary_window.reset_at,
|
|
}
|
|
if usage.primary_window
|
|
else None,
|
|
"secondary_window": {
|
|
"used_percent": usage.secondary_window.used_percent,
|
|
"reset_at": usage.secondary_window.reset_at,
|
|
}
|
|
if usage.secondary_window
|
|
else None,
|
|
},
|
|
}
|
|
|
|
async def validate_token(self, access_token: str) -> bool:
|
|
self.validated_tokens.append(access_token)
|
|
return access_token not in self.invalid_tokens
|
|
|
|
|
|
def make_usage(primary: int, secondary: int = 0, *, checked_at: int | None = None):
|
|
return UsageSnapshot(
|
|
checked_at=checked_at or int(time.time()),
|
|
primary_window=UsageWindow(used_percent=primary, reset_at=int(time.time()) + 300),
|
|
secondary_window=UsageWindow(
|
|
used_percent=secondary, reset_at=int(time.time()) + 300
|
|
),
|
|
)
|
|
|
|
|
|
def make_account(
|
|
email: str,
|
|
*,
|
|
token: str,
|
|
refresh_token: str = "refresh",
|
|
token_refresh_at: int | None = None,
|
|
usage: UsageSnapshot | None = None,
|
|
disabled: bool = False,
|
|
) -> AccountRecord:
|
|
return AccountRecord(
|
|
email=email,
|
|
access_token=token,
|
|
refresh_token=refresh_token,
|
|
token_refresh_at=token_refresh_at or int(time.time()) + 600,
|
|
usage=usage,
|
|
usage_checked_at=usage.checked_at if usage is not None else None,
|
|
disabled=disabled,
|
|
)
|
|
|
|
|
|
def make_store(tmp_path: Path, state: StateFile) -> JsonStateStore:
|
|
store = JsonStateStore(tmp_path / "accounts.json")
|
|
store.save(state)
|
|
return store
|
|
|
|
|
|
def make_manager(
|
|
store: JsonStateStore,
|
|
client: FakeClient,
|
|
*,
|
|
threshold: int = 95,
|
|
stale_seconds: int = 3600,
|
|
) -> AccountManager:
|
|
return AccountManager(
|
|
store,
|
|
client,
|
|
Settings(
|
|
data_dir=store.path.parent,
|
|
exhausted_usage_threshold=threshold,
|
|
usage_stale_seconds=stale_seconds,
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prefers_active_account_when_usable(tmp_path: Path) -> None:
|
|
active = make_account("a@example.com", token="tok-a", usage=make_usage(20, 0))
|
|
second = make_account("b@example.com", token="tok-b", usage=make_usage(80, 0))
|
|
store = make_store(
|
|
tmp_path,
|
|
StateFile(active_account="a@example.com", accounts=[active, second]),
|
|
)
|
|
client = FakeClient()
|
|
|
|
payload = await make_manager(store, client).issue_token_response()
|
|
|
|
assert payload["token"] == "tok-a"
|
|
assert client.fetched_usage_tokens == []
|
|
assert client.validated_tokens == ["tok-a"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refreshes_stale_active_usage_before_deciding(tmp_path: Path) -> None:
|
|
stale = int(time.time()) - 7200
|
|
active = make_account("a@example.com", token="tok-a", usage=make_usage(20, 0, checked_at=stale))
|
|
second = make_account("b@example.com", token="tok-b", usage=make_usage(80, 0))
|
|
store = make_store(
|
|
tmp_path,
|
|
StateFile(active_account="a@example.com", accounts=[active, second]),
|
|
)
|
|
client = FakeClient(usage_by_token={"tok-a": make_usage(21, 0)})
|
|
|
|
payload = await make_manager(store, client).issue_token_response()
|
|
|
|
assert payload["token"] == "tok-a"
|
|
assert client.fetched_usage_tokens == ["tok-a"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_falls_back_to_highest_primary_usage_when_active_unusable(tmp_path: Path) -> None:
|
|
active = make_account("a@example.com", token="tok-a", usage=make_usage(95, 0))
|
|
low = make_account("b@example.com", token="tok-b", usage=make_usage(40, 0))
|
|
high = make_account("c@example.com", token="tok-c", usage=make_usage(70, 0))
|
|
store = make_store(
|
|
tmp_path,
|
|
StateFile(active_account="a@example.com", accounts=[active, low, high]),
|
|
)
|
|
client = FakeClient()
|
|
|
|
payload = await make_manager(store, client).issue_token_response()
|
|
state = store.load()
|
|
|
|
assert payload["token"] == "tok-c"
|
|
assert state.active_account == "c@example.com"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_skips_disabled_accounts(tmp_path: Path) -> None:
|
|
active = make_account("a@example.com", token="tok-a", usage=make_usage(20, 0), disabled=True)
|
|
second = make_account("b@example.com", token="tok-b", usage=make_usage(70, 0))
|
|
store = make_store(
|
|
tmp_path,
|
|
StateFile(active_account="a@example.com", accounts=[active, second]),
|
|
)
|
|
client = FakeClient()
|
|
|
|
payload = await make_manager(store, client).issue_token_response()
|
|
|
|
assert payload["token"] == "tok-b"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_secondary_100_makes_account_unusable(tmp_path: Path) -> None:
|
|
active = make_account("a@example.com", token="tok-a", usage=make_usage(20, 100))
|
|
second = make_account("b@example.com", token="tok-b", usage=make_usage(30, 0))
|
|
store = make_store(
|
|
tmp_path,
|
|
StateFile(active_account="a@example.com", accounts=[active, second]),
|
|
)
|
|
client = FakeClient()
|
|
|
|
payload = await make_manager(store, client).issue_token_response()
|
|
|
|
assert payload["token"] == "tok-b"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refreshes_token_before_validation(tmp_path: Path) -> None:
|
|
account = make_account(
|
|
"a@example.com",
|
|
token="old-token",
|
|
refresh_token="ref-a",
|
|
token_refresh_at=int(time.time()) - 1,
|
|
usage=make_usage(20, 0),
|
|
)
|
|
store = make_store(tmp_path, StateFile(active_account="a@example.com", accounts=[account]))
|
|
client = FakeClient(refresh_map={"ref-a": ("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-a"]
|
|
assert client.validated_tokens == ["new-token"]
|
|
assert saved.accounts[0].access_token == "new-token"
|
|
assert saved.accounts[0].refresh_token == "new-refresh"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_invalid_token_moves_account_to_failed_json(tmp_path: Path) -> None:
|
|
bad = make_account("bad@example.com", token="tok-bad", usage=make_usage(20, 0))
|
|
good = make_account("good@example.com", token="tok-good", usage=make_usage(30, 0))
|
|
store = make_store(tmp_path, StateFile(active_account="bad@example.com", accounts=[bad, good]))
|
|
client = FakeClient(invalid_tokens={"tok-bad"})
|
|
|
|
payload = await make_manager(store, client).issue_token_response()
|
|
state = store.load()
|
|
failed = json.loads((tmp_path / "failed.json").read_text())
|
|
|
|
assert payload["token"] == "tok-good"
|
|
assert [account.email for account in state.accounts] == ["good@example.com"]
|
|
assert failed["accounts"][0]["email"] == "bad@example.com"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rereads_disk_between_requests(tmp_path: Path) -> None:
|
|
first = make_account("a@example.com", token="tok-a", usage=make_usage(20, 0))
|
|
store = make_store(tmp_path, StateFile(active_account="a@example.com", accounts=[first]))
|
|
client = FakeClient()
|
|
manager = make_manager(store, client)
|
|
|
|
first_payload = await manager.issue_token_response()
|
|
assert first_payload["token"] == "tok-a"
|
|
|
|
replacement = make_account("b@example.com", token="tok-b", usage=make_usage(10, 0))
|
|
store.save(StateFile(active_account="b@example.com", accounts=[replacement]))
|
|
|
|
second_payload = await manager.issue_token_response()
|
|
assert second_payload["token"] == "tok-b"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_raises_when_no_usable_accounts(tmp_path: Path) -> None:
|
|
disabled = make_account("a@example.com", token="tok-a", usage=make_usage(10, 0), disabled=True)
|
|
exhausted = make_account("b@example.com", token="tok-b", usage=make_usage(95, 0))
|
|
store = make_store(tmp_path, StateFile(accounts=[disabled, exhausted]))
|
|
client = FakeClient()
|
|
|
|
with pytest.raises(NoUsableAccountError):
|
|
await make_manager(store, client).issue_token_response()
|