200 lines
7.7 KiB
Python
200 lines
7.7 KiB
Python
"""Tests for cloud.models (fallback routing) and cloud.cost.
|
|
|
|
No real model turns, no network, no secrets: usage clients are fakes
|
|
and the CostStore uses an in-memory SQLite database.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
|
|
# cloud/ has no __init__.py by design (the coordinator creates it), so
|
|
# make the project root importable regardless of the discover top level.
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(
|
|
os.path.dirname(os.path.abspath(__file__)))))
|
|
|
|
from cloud.cost import CostStore, estimate_tokens, fetch_account_usage
|
|
from cloud.models import FALLBACKABLE, ModelRouter, is_fallbackable, record_fallback
|
|
|
|
|
|
class FakeError(Exception):
|
|
def __init__(self, status=None, code=None):
|
|
super().__init__(status or code)
|
|
self.status = status
|
|
self.code = code
|
|
|
|
|
|
class TestIsFallbackable(unittest.TestCase):
|
|
def test_fallbackable_status_codes(self):
|
|
for code in (404, 429, 500, 502, 503, 504):
|
|
self.assertTrue(is_fallbackable(code), code)
|
|
self.assertIn(404, FALLBACKABLE)
|
|
self.assertIn(429, FALLBACKABLE)
|
|
self.assertIn("timeout", FALLBACKABLE)
|
|
|
|
def test_timeout_string(self):
|
|
self.assertTrue(is_fallbackable("timeout"))
|
|
|
|
def test_not_fallbackable(self):
|
|
for code in (400, 401, 403, 200):
|
|
self.assertFalse(is_fallbackable(code), code)
|
|
self.assertFalse(is_fallbackable("other"))
|
|
|
|
def test_object_with_status_or_code(self):
|
|
self.assertTrue(is_fallbackable(FakeError(status=429)))
|
|
self.assertTrue(is_fallbackable(FakeError(code=503)))
|
|
self.assertTrue(is_fallbackable(FakeError(status="timeout")))
|
|
self.assertFalse(is_fallbackable(FakeError(status=401)))
|
|
self.assertFalse(is_fallbackable(FakeError(code=400)))
|
|
|
|
|
|
class TestModelRouter(unittest.TestCase):
|
|
def setUp(self):
|
|
self.router = ModelRouter(["gpt-6.1-sol", "gpt-6-sol"])
|
|
|
|
def test_empty_priority_list_rejected(self):
|
|
with self.assertRaises(ValueError):
|
|
ModelRouter([])
|
|
|
|
def test_first(self):
|
|
self.assertEqual(self.router.first(), "gpt-6.1-sol")
|
|
|
|
def test_next_after_in_priority_order(self):
|
|
self.assertEqual(
|
|
self.router.next_after("gpt-6.1-sol", 404), "gpt-6-sol")
|
|
self.assertEqual(
|
|
self.router.next_after("gpt-6.1-sol", "timeout"), "gpt-6-sol")
|
|
|
|
def test_next_after_exhausted_returns_none(self):
|
|
self.assertIsNone(self.router.next_after("gpt-6-sol", 500))
|
|
self.assertIsNone(self.router.next_after("gpt-6-sol", 429))
|
|
|
|
def test_next_after_non_fallbackable_returns_none(self):
|
|
self.assertIsNone(self.router.next_after("gpt-6.1-sol", 401))
|
|
self.assertIsNone(self.router.next_after("gpt-6.1-sol", 400))
|
|
|
|
def test_route_walks_attempts(self):
|
|
self.assertEqual(
|
|
self.router.route([("gpt-6.1-sol", 404)]), "gpt-6-sol")
|
|
self.assertIsNone(self.router.route(
|
|
[("gpt-6.1-sol", 404), ("gpt-6-sol", 500)]))
|
|
self.assertIsNone(self.router.route([("gpt-6.1-sol", 401)]))
|
|
self.assertEqual(self.router.route([]), "gpt-6.1-sol")
|
|
|
|
|
|
class TestRecordFallback(unittest.TestCase):
|
|
def test_event_recorded_on_task_record(self):
|
|
task = {"id": "task-1"}
|
|
event = record_fallback(task, "gpt-6.1-sol", "gpt-6-sol", "429")
|
|
self.assertEqual(event, {
|
|
"type": "model_fallback",
|
|
"task_id": "task-1",
|
|
"from_model": "gpt-6.1-sol",
|
|
"to_model": "gpt-6-sol",
|
|
"reason": "429",
|
|
})
|
|
self.assertIn(event, task["events"])
|
|
|
|
def test_event_appended_to_existing_events(self):
|
|
task = {"id": "task-2", "events": [{"type": "started"}]}
|
|
record_fallback(task, "a", "b", "500")
|
|
self.assertEqual(len(task["events"]), 2)
|
|
self.assertEqual(task["events"][1]["type"], "model_fallback")
|
|
|
|
|
|
class TestFetchAccountUsage(unittest.TestCase):
|
|
def test_success_returns_available_true(self):
|
|
data = {"balance": 123, "currency": "USD"}
|
|
result = fetch_account_usage(lambda: data)
|
|
self.assertEqual(result, {"available": True, "data": data})
|
|
|
|
def test_raising_client_returns_unavailable_not_raise(self):
|
|
def boom():
|
|
raise RuntimeError("no route")
|
|
self.assertEqual(
|
|
fetch_account_usage(boom), {"available": False, "data": None})
|
|
|
|
def test_none_and_empty_dict_return_unavailable(self):
|
|
self.assertEqual(
|
|
fetch_account_usage(lambda: None),
|
|
{"available": False, "data": None})
|
|
self.assertEqual(
|
|
fetch_account_usage(lambda: {}),
|
|
{"available": False, "data": None})
|
|
|
|
|
|
class TestEstimateTokens(unittest.TestCase):
|
|
def test_chars_divided_by_four(self):
|
|
self.assertEqual(estimate_tokens("abcd", "efgh"), 2)
|
|
self.assertEqual(estimate_tokens("a" * 10, "b" * 7), 4) # 17 // 4
|
|
self.assertEqual(estimate_tokens("", ""), 0)
|
|
|
|
|
|
class TestCostStore(unittest.TestCase):
|
|
def setUp(self):
|
|
self.store = CostStore(":memory:")
|
|
|
|
def test_record_computes_total_and_day(self):
|
|
row_id = self.store.record(
|
|
"t1", "p1", "gpt-6.1-sol", input_tokens=100, output_tokens=50,
|
|
usd_micros=2000, source="actual", fallback_count=1)
|
|
self.assertIsInstance(row_id, int)
|
|
report = self.store.report_by_task("t1")
|
|
self.assertEqual(report["total_tokens"], 150)
|
|
self.assertEqual(report["usd_micros"], 2000)
|
|
self.assertEqual(report["turns"], 1)
|
|
self.assertEqual(report["fallback_count"], 1)
|
|
self.assertEqual(report["actual_tokens"], 150)
|
|
self.assertEqual(report["estimated_tokens"], 0)
|
|
|
|
def test_invalid_source_rejected(self):
|
|
with self.assertRaises(ValueError):
|
|
self.store.record("t1", "p1", "m", source="guess")
|
|
|
|
def test_actual_and_estimated_kept_separate(self):
|
|
self.store.record("t1", "p1", "m", input_tokens=100,
|
|
output_tokens=0, source="actual")
|
|
self.store.record_estimated("t1", "p1", "m",
|
|
"x" * 40, "y" * 40) # 80 // 4 = 20
|
|
report = self.store.report_by_task("t1")
|
|
self.assertEqual(report["total_tokens"], 120)
|
|
self.assertEqual(report["actual_tokens"], 100)
|
|
self.assertEqual(report["estimated_tokens"], 20)
|
|
self.assertEqual(report["turns"], 2)
|
|
|
|
def test_record_estimated_shape(self):
|
|
self.store.record_estimated("t9", "p9", "m", "abcd", "efgh")
|
|
report = self.store.report_by_task("t9")
|
|
self.assertEqual(report["estimated_tokens"], 2)
|
|
self.assertEqual(report["actual_tokens"], 0)
|
|
self.assertEqual(report["usd_micros"], 0)
|
|
|
|
def test_report_by_project_and_day(self):
|
|
self.store.record("t1", "p1", "m", input_tokens=10, output_tokens=0)
|
|
self.store.record("t2", "p1", "m", input_tokens=30, output_tokens=0)
|
|
self.store.record("t3", "p2", "m", input_tokens=999, output_tokens=0)
|
|
project = self.store.report_by_project("p1")
|
|
self.assertEqual(project["total_tokens"], 40)
|
|
self.assertEqual(project["turns"], 2)
|
|
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
day = self.store.report_by_day(today)
|
|
self.assertEqual(day["total_tokens"], 10 + 30 + 999)
|
|
self.assertEqual(day["turns"], 3)
|
|
self.assertEqual(self.store.report_by_day("2000-01-01")["turns"], 0)
|
|
|
|
def test_empty_report_is_zeroed(self):
|
|
report = self.store.report_by_task("missing")
|
|
self.assertEqual(report, {
|
|
"total_tokens": 0,
|
|
"usd_micros": 0,
|
|
"turns": 0,
|
|
"fallback_count": 0,
|
|
"actual_tokens": 0,
|
|
"estimated_tokens": 0,
|
|
})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|