69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
"""Tests for migration 0007_usage — token-budget accounting."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from claire import events as ev
|
|
from claire.db import migrate, open_db
|
|
from claire.hlc import HLCGenerator
|
|
from claire.web import service
|
|
|
|
|
|
def _setup() -> tuple:
|
|
conn = open_db(":memory:")
|
|
migrate(conn)
|
|
return conn, HLCGenerator("test-machine")
|
|
|
|
|
|
def test_usage_recorded_projects() -> None:
|
|
conn, gen = _setup()
|
|
ev.append(conn, gen, ev.UsageRecorded(
|
|
source="nl", model="haiku",
|
|
input_tokens=120, output_tokens=40, cache_read_tokens=2000,
|
|
))
|
|
row = conn.execute("SELECT * FROM usage").fetchone()
|
|
assert row["source"] == "nl"
|
|
assert row["input_tokens"] == 120
|
|
assert row["output_tokens"] == 40
|
|
assert row["cache_read_tokens"] == 2000
|
|
assert len(row["day"]) == 10 # YYYY-MM-DD
|
|
|
|
|
|
def test_record_usage_service() -> None:
|
|
conn, gen = _setup()
|
|
service.record_usage(conn, gen, source="triage", model="haiku",
|
|
input_tokens=50, output_tokens=10)
|
|
service.record_usage(conn, gen, source="nl", model="haiku",
|
|
input_tokens=80, output_tokens=20)
|
|
from claire import read
|
|
roll = read.usage_rollup(conn)
|
|
assert roll["total_tokens"] == 160
|
|
assert roll["calls"] == 2
|
|
sources = {s["source"] for s in roll["by_source"]}
|
|
assert sources == {"triage", "nl"}
|
|
|
|
|
|
def test_budget_status_no_cap() -> None:
|
|
conn, gen = _setup()
|
|
service.record_usage(conn, gen, source="nl", model="haiku",
|
|
input_tokens=100, output_tokens=100)
|
|
b = service.budget_status(conn)
|
|
assert b["used_tokens"] == 200
|
|
assert b["daily_token_cap"] == 0 # default — unlimited
|
|
assert b["over_cap"] is False
|
|
assert b["over_low_priority_floor"] is False
|
|
|
|
|
|
def test_usage_replay_roundtrip() -> None:
|
|
conn, gen = _setup()
|
|
for i in range(5):
|
|
service.record_usage(conn, gen, source="nl", model="haiku",
|
|
input_tokens=10 * i, output_tokens=i)
|
|
before = conn.execute(
|
|
"SELECT COALESCE(SUM(input_tokens+output_tokens),0) FROM usage"
|
|
).fetchone()[0]
|
|
ev.replay(conn)
|
|
after = conn.execute(
|
|
"SELECT COALESCE(SUM(input_tokens+output_tokens),0) FROM usage"
|
|
).fetchone()[0]
|
|
assert before == after
|
|
assert conn.execute("SELECT COUNT(*) FROM usage").fetchone()[0] == 5
|