"""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 from unittest import mock # 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, price_usd_micros) 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, completions=7) 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) self.assertEqual(report["llm_requests"], 7) def test_llm_requests_sums_completions_not_turns(self): self.store.record("t1", "p1", "m", input_tokens=10, output_tokens=0, source="actual", completions=3) self.store.record("t1", "p1", "m", input_tokens=20, output_tokens=0, source="actual", completions=5) report = self.store.report_by_task("t1") self.assertEqual(report["turns"], 2) self.assertEqual(report["llm_requests"], 8) 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, "cached_input_tokens": 0, "unpriced_actual_tokens": 0, "llm_requests": 0, }) class TestModelPrices(unittest.TestCase): def test_list_prices(self): # 2026-10-26 中继实测 5 调用回合：(58243-45960)*2 # + 45960*0.10 + 1138*10 = 40542 微美元。 self.assertEqual( price_usd_micros("gpt-6.1-sol", 58243, 45960, 1138), 40542) self.assertEqual( price_usd_micros("gpt-6-sol", 1000, 200, 100), int(round((800 * 2.00 + 200 * 0.20 + 100 * 10.00)))) def test_unknown_model_unpriced(self): self.assertIsNone(price_usd_micros("gpt-5.6-sol", 100, 0, 100)) self.assertIsNone(price_usd_micros(None, 100, 0, 100)) def test_cached_clamped_to_input(self): self.assertEqual( price_usd_micros("gpt-6.1-sol", 100, 500, 100), price_usd_micros("gpt-6.1-sol", 100, 100, 100)) def test_env_override(self): with mock.patch.dict(os.environ, { "CLOUD_MODEL_PRICES_JSON": '{"gpt-6.1-sol": [4.00, 0.20, 20.00]}'}): self.assertEqual( price_usd_micros("gpt-6.1-sol", 1000, 0, 0), 4000) with mock.patch.dict(os.environ, { "CLOUD_MODEL_PRICES_JSON": 'not json'}): # 坏覆盖不炸，回落默认表。 self.assertEqual( price_usd_micros("gpt-6.1-sol", 1000, 0, 0), 2000) class TestCachedColumn(unittest.TestCase): def test_record_keeps_cached_and_unpriced(self): store = CostStore(":memory:") store.record("t1", "p1", "gpt-6.1-sol", input_tokens=1000, output_tokens=100, cached_input_tokens=400, usd_micros=2240, source="actual") store.record("t2", "p1", "gpt-5.6-sol", input_tokens=500, output_tokens=50, cached_input_tokens=0, usd_micros=None, source="actual") report = store.report_by_project("p1") self.assertEqual(report["cached_input_tokens"], 400) self.assertEqual(report["unpriced_actual_tokens"], 550) self.assertEqual(report["usd_micros"], 2240) def test_migration_adds_cached_column_to_old_table(self): import sqlite3 import tempfile path = os.path.join(tempfile.mkdtemp(), "c.db") conn = sqlite3.connect(path) conn.execute( "CREATE TABLE cost_records (id INTEGER PRIMARY KEY" " AUTOINCREMENT, task_id TEXT, project_id TEXT, model TEXT," " input_tokens INTEGER, output_tokens INTEGER," " total_tokens INTEGER, usd_micros INTEGER, source TEXT," " fallback_count INTEGER DEFAULT 0, day TEXT," " created_at REAL)") conn.commit() conn.close() store = CostStore(path) cols = {row[1] for row in store._conn.execute( "PRAGMA table_info(cost_records)").fetchall()} self.assertIn("cached_input_tokens", cols) store.record("t1", "p1", "m", input_tokens=10, output_tokens=5, cached_input_tokens=3, source="actual") self.assertEqual(store.report_by_task("t1")["cached_input_tokens"], 3) if __name__ == "__main__": unittest.main()