ec7851a819
WI-2 of quota-multi-provider-redesign. - total_credits_usd REAL + total_usage_usd REAL columns added to provider_balance_observations table in schema.sql + idempotent ALTER. - _BALANCE_PARSERS now returns 3-tuple (balance, total_credits, total_usage). OpenRouter parser extracts pool size from /credits response; non-pool providers return (x, None, None) for backward compat. - record_balance persists all three columns. - Updated test_provider_balance.py column-list pins and round-trip assertion.
93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
"""Tests for the provider_balance_observations table and poller helpers.
|
|
|
|
All tests seed a throwaway SQLite DB from schema.sql and assert on queried
|
|
rows, not on stdout strings.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
import poller
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
SCHEMA_SQL = (ROOT / "config" / "schema.sql").read_text()
|
|
|
|
|
|
def _make_db() -> sqlite3.Connection:
|
|
conn = sqlite3.connect(":memory:")
|
|
conn.execute("PRAGMA foreign_keys = ON")
|
|
conn.executescript(SCHEMA_SQL)
|
|
return conn
|
|
|
|
|
|
def _table_columns(conn: sqlite3.Connection, table: str) -> list[str]:
|
|
return [row[1] for row in conn.execute(f"PRAGMA table_info({table})")]
|
|
|
|
|
|
def test_provider_balance_observations_schema() -> None:
|
|
"""QA (a): schema.sql creates the expected columns."""
|
|
conn = _make_db()
|
|
assert _table_columns(conn, "provider_balance_observations") == [
|
|
"id",
|
|
"provider",
|
|
"balance_usd",
|
|
"total_credits_usd",
|
|
"total_usage_usd",
|
|
"observed_at",
|
|
]
|
|
|
|
|
|
def test_record_balance_round_trip() -> None:
|
|
"""QA (b): record_balance inserts UTC ISO8601 with correct values."""
|
|
conn = _make_db()
|
|
before = datetime.now(timezone.utc)
|
|
poller.record_balance(conn, "openrouter", 42.5, 50.0, 7.5)
|
|
after = datetime.now(timezone.utc)
|
|
|
|
row = conn.execute(
|
|
"SELECT provider, balance_usd, total_credits_usd, total_usage_usd, observed_at "
|
|
"FROM provider_balance_observations"
|
|
).fetchone()
|
|
assert row is not None
|
|
assert row[0] == "openrouter"
|
|
assert row[1] == pytest.approx(42.5)
|
|
assert row[2] == pytest.approx(50.0)
|
|
assert row[3] == pytest.approx(7.5)
|
|
observed = datetime.fromisoformat(row[4])
|
|
assert observed.tzinfo is not None
|
|
assert before <= observed <= after
|
|
|
|
|
|
def test_ensure_provider_balance_table_idempotent() -> None:
|
|
"""QA (c): _ensure_provider_balance_table runs twice on an empty DB."""
|
|
conn = sqlite3.connect(":memory:")
|
|
conn.execute("PRAGMA foreign_keys = ON")
|
|
poller._ensure_provider_balance_table(conn)
|
|
poller._ensure_provider_balance_table(conn)
|
|
assert _table_columns(conn, "provider_balance_observations") == [
|
|
"id",
|
|
"provider",
|
|
"balance_usd",
|
|
"total_credits_usd",
|
|
"total_usage_usd",
|
|
"observed_at",
|
|
]
|
|
idx = conn.execute(
|
|
"SELECT name FROM sqlite_master WHERE type='index' AND name='idx_provider_balance'"
|
|
).fetchone()
|
|
assert idx is not None
|
|
|
|
|
|
def test_record_balance_fails_without_table() -> None:
|
|
"""QA (d): no hidden auto-create; missing table raises OperationalError."""
|
|
conn = sqlite3.connect(":memory:")
|
|
conn.execute("PRAGMA foreign_keys = ON")
|
|
# Table intentionally not created.
|
|
with pytest.raises(sqlite3.OperationalError):
|
|
poller.record_balance(conn, "openrouter", 1.0)
|