cloud: production control plane (API, auth/RBAC, queue, sandbox, memory/KB, fallback/cost, webhooks)
This commit is contained in:
host committed
1 parent
b97c34465f
commit
b74d99e959
17 files changed
+3420
No files matched your search
@@ -0,0 +1,199 @@
|
||||
"""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()
|
||||
Reference in new issue
Block a user