shujakuin/shusub2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: show Pricing Monitor source balances

Browse Source
This commit is contained in:
yunyaozhou
2026-08-03 01:08:49 +08:00
parent 845f103bd1
commit 49b2f963b9
5 changed files with 460 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_payload.py
+133 -5
View File
@@ -39,8 +39,8 @@ class Sub2APIQuotaTUITests(unittest.TestCase):
"today_tokens": 2000,
"today_requests": 4,
"windows": [
{"id": "five-hour", "used_percent": 80, "remaining_percent": 20, "remaining_balance_usd": 4, "reset": "2026-06-09T15:00:00+08:00"},
{"id": "weekly", "used_percent": 10, "remaining_percent": 90, "remaining_balance_usd": 108, "reset": "2026-06-10T15:00:00+08:00"},
{"id": "five-hour", "used_percent": 80, "remaining_percent": 20, "reset": "2026-06-09T15:00:00+08:00"},
{"id": "weekly", "used_percent": 10, "remaining_percent": 90, "reset": "2026-06-10T15:00:00+08:00"},
],
},
{
@@ -72,8 +72,8 @@ class Sub2APIQuotaTUITests(unittest.TestCase):
self.assertEqual(rows[0]["kind_label"], "usage")
self.assertEqual(rows[1]["five_hour"], "80%/20%")
self.assertEqual(rows[1]["weekly"], "10%/90%")
self.assertEqual(rows[1]["five_hour_balance"], "$4")
self.assertEqual(rows[1]["weekly_balance"], "$108")
self.assertNotIn("five_hour_balance", rows[1])
self.assertNotIn("weekly_balance", rows[1])
self.assertIn("today $0.75", mod.summary_line(payload))
def test_filter_matches_kind_and_name(self) -> None:
@@ -90,6 +90,89 @@ class Sub2APIQuotaTUITests(unittest.TestCase):
self.assertEqual([row["name"] for row in mod.normalize_account_rows(payload, "sfast")], ["beta"])
self.assertEqual([row["name"] for row in mod.normalize_account_rows(payload, "anthropic")], ["alpha"])
def test_normalize_pricing_rows_uses_pricing_monitor_source_balances(self) -> None:
mod = load_module()
payload = {
"generated_at": "2026-08-02T12:00:00Z",
"sources": [
{
"name": "code-plan",
"source_kind": "newapi",
"health_state": "healthy",
"last_success_at": "2026-08-02T11:59:00Z",
"balance_available": True,
"balance": {"available": 5000000, "available_cny": 10, "unit": "quota"},
},
{
"name": "kedaya",
"source_kind": "remote_monitor",
"health_state": "error",
"last_error": "upstream request failed",
"balance_available": False,
"balance": {"available": 99, "available_cny": 99, "unit": "balance"},
},
],
}
rows = mod.normalize_pricing_rows(payload)
self.assertEqual([row["name"] for row in rows], ["code-plan", "kedaya"])
self.assertEqual(rows[0]["balance"], 5000000)
self.assertEqual(rows[0]["balance_cny"], 10)
self.assertEqual(rows[0]["unit"], "quota")
self.assertEqual(rows[0]["status"], "healthy")
self.assertEqual(rows[1]["status"], "error")
self.assertEqual([row["name"] for row in mod.normalize_pricing_rows(payload, "keday")], ["kedaya"])
self.assertEqual(mod.pricing_summary(payload, rows), "upstreams 1/2 healthy | CNY ¥10")
out = io.StringIO()
with contextlib.redirect_stdout(out):
mod.print_pricing_once(payload)
self.assertIn("code-plan", out.getvalue())
self.assertIn("5,000,000", out.getvalue())
self.assertIn("¥10", out.getvalue())
def test_default_pricing_url_reads_env_or_config_file(self) -> None:
mod = load_module()
with tempfile.TemporaryDirectory() as tmp:
config_file = Path(tmp) / "shusub2" / "pricing-url"
config_file.parent.mkdir(parents=True)
config_file.write_text("https://pricing.example.test/api/ui-data?view=accounts\n", encoding="utf-8")
old_url = os.environ.get("SHUSUB2_PRICING_URL")
old_file = os.environ.get("SHUSUB2_PRICING_URL_FILE")
os.environ["SHUSUB2_PRICING_URL_FILE"] = str(config_file)
os.environ.pop("SHUSUB2_PRICING_URL", None)
try:
self.assertEqual(
mod.default_pricing_url(),
"https://pricing.example.test/api/ui-data?view=accounts",
)
os.environ["SHUSUB2_PRICING_URL"] = "https://env.example.test/api/ui-data?view=accounts"
self.assertEqual(
mod.default_pricing_url(),
"https://env.example.test/api/ui-data?view=accounts",
)
finally:
if old_url is None:
os.environ.pop("SHUSUB2_PRICING_URL", None)
else:
os.environ["SHUSUB2_PRICING_URL"] = old_url
if old_file is None:
os.environ.pop("SHUSUB2_PRICING_URL_FILE", None)
else:
os.environ["SHUSUB2_PRICING_URL_FILE"] = old_file
def test_fetch_pricing_payload_requires_sources_list(self) -> None:
mod = load_module()
original_fetch = mod.fetch_payload
mod.fetch_payload = lambda *args, **kwargs: {"view": "accounts", "sources": []}
try:
self.assertEqual(mod.fetch_pricing_payload("https://price.example/api/ui-data?view=accounts", 1)["sources"], [])
mod.fetch_payload = lambda *args, **kwargs: {"view": "accounts"}
with self.assertRaisesRegex(RuntimeError, "source balances"):
mod.fetch_pricing_payload("https://price.example/api/ui-data?view=accounts", 1)
finally:
mod.fetch_payload = original_fetch
def test_infers_public_status_url_from_accounts_url(self) -> None:
mod = load_module()
@@ -252,7 +335,7 @@ class Sub2APIQuotaTUITests(unittest.TestCase):
self.assertEqual(mod.monitor_availability(rows[0], status_payload), "ok")
self.assertIn("availability", out.getvalue())
self.assertIn("5h bal", out.getvalue())
self.assertNotIn("5h bal", out.getvalue())
self.assertIn(" ok\n", out.getvalue())
def test_monitor_availability_is_dash_without_bound_monitor(self) -> None:
@@ -756,6 +839,19 @@ class DashboardLayoutTests(unittest.IsolatedAsyncioTestCase):
],
}
mod.fetch_optional_payload = lambda *args, **kwargs: ({}, "")
mod.fetch_pricing_payload = lambda *args, **kwargs: {
"generated_at": "2026-07-24T12:00:00+08:00",
"sources": [
{
"name": "code-plan",
"source_kind": "newapi",
"health_state": "healthy",
"last_success_at": "2026-07-24T12:00:00+08:00",
"balance_available": True,
"balance": {"available": 5000000, "available_cny": 10, "unit": "quota"},
}
],
}
mod.fetch_key_usage_payload = lambda *args, **kwargs: {
"date": "2026-07-24",
"trend": [{"api_key_id": 1, "key_name": "wmy", "requests": 3, "tokens": 300}],
@@ -805,6 +901,7 @@ class DashboardLayoutTests(unittest.IsolatedAsyncioTestCase):
rc = mod.run_textual(
"https://example/accounts",
"",
"https://example/pricing?view=accounts",
"https://example/usage",
"token",
"https://example/cn",
@@ -846,9 +943,40 @@ class DashboardLayoutTests(unittest.IsolatedAsyncioTestCase):
await pilot.press(key)
await pilot.pause()
self.assertEqual(screen.focused.id, expected_id)
await pilot.press("p")
await pilot.pause()
self.assertEqual(type(app.screen).__name__, "PricingScreen")
self.assertEqual(app.screen.focused.id, "pricing")
self.assertEqual(app.screen.rows[0]["name"], "code-plan")
class PageSelectionTests(unittest.TestCase):
def test_once_pricing_prints_pricing_monitor_sources(self) -> None:
mod = load_module()
original_fetch = mod.fetch_pricing_payload
mod.fetch_pricing_payload = lambda *args, **kwargs: {
"generated_at": "2026-08-02T12:00:00Z",
"sources": [
{
"name": "code-plan",
"source_kind": "newapi",
"health_state": "healthy",
"last_success_at": "2026-08-02T11:59:00Z",
"balance": {"available": 5000000, "available_cny": 10, "unit": "quota"},
}
],
}
out = io.StringIO()
try:
with contextlib.redirect_stdout(out):
rc = mod.main(["--once", "--pricing", "--no-version-check"])
finally:
mod.fetch_pricing_payload = original_fetch
self.assertEqual(rc, 0)
self.assertIn("code-plan", out.getvalue())
self.assertIn("¥10", out.getvalue())
def test_dedicated_page_flags_are_mutually_exclusive(self) -> None:
mod = load_module()
old_token = os.environ.pop("SHUSUB2_LOGS_TOKEN", None)