shujakuin/shusub2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
2d3d433f4f8ad71494211fed286491e6e2fa25c7
shusub2/tests/test_payload.py
T
shujakuin 2d3d433f4f feat: show first-token latency, decode tok/s, per-key colors 
- logs page adds First (first_token_ms) and Tok/s columns; Tok/s is
  output_tokens / (latency - first_token) over the decode window
- detail line shows the same rate after first-token latency
- API key names render in a stable per-key color (crc32 over a 10-color
  palette) so rows from the same key group visually; --once stays plain
- --once --logs prints the new first/tok-s columns

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-21 11:19:06 +08:00

550 lines
23 KiB
Python
Raw Blame History

from __future__ import annotations
import contextlib
import io
import os
from pathlib import Path
import importlib.util
import sys
import tempfile
import unittest
def load_module():
module_path = Path(__file__).resolve().parents[1] / "sub2api_quota_tui.py"
spec = importlib.util.spec_from_file_location("sub2api_quota_tui", module_path)
assert spec and spec.loader
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
class Sub2APIQuotaTUITests(unittest.TestCase):
def test_normalize_rows_sorts_by_usage_and_formats_windows(self) -> None:
mod = load_module()
payload = {
"generated_at": "2026-06-09T12:00:00+08:00",
"totals": {"total_accounts": 2, "usable_accounts": 2, "today_cost_usd": 0.75, "today_tokens": 3000, "today_requests": 9},
"accounts": [
{
"id": 1,
"name": "quota-low",
"routing_group": "slow",
"provider": "anthropic",
"kind": "quota_limited",
"status": "ok",
"account_type": "oauth",
"today_cost_usd": 0.25,
"today_tokens": 2000,
"today_requests": 4,
"windows": [
{"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"},
],
},
{
"id": 2,
"name": "payg-high",
"routing_group": "fast",
"provider": "openai",
"kind": "pay_as_you_go",
"status": "ok",
"account_type": "apikey",
"daily_quota_limit": 100,
"daily_quota_used": 91.2,
"daily_quota_remaining": 8.8,
"daily_quota_used_percent": 91.2,
"today_cost_usd": 0.5,
"today_tokens": 1000,
"today_requests": 5,
"windows": [],
},
],
}
rows = mod.normalize_account_rows(payload)
self.assertEqual([row["name"] for row in rows], ["payg-high", "quota-low"])
self.assertEqual(rows[0]["routing_group"], "fast")
self.assertEqual(rows[0]["provider"], "openai")
self.assertEqual(rows[0]["daily_quota_cell"], "91.2/100")
self.assertEqual(rows[0]["kind_label"], "usage")
self.assertEqual(rows[1]["five_hour"], "80%/20%")
self.assertEqual(rows[1]["weekly"], "10%/90%")
self.assertIn("today $0.75", mod.summary_line(payload))
def test_filter_matches_kind_and_name(self) -> None:
mod = load_module()
payload = {
"accounts": [
{"id": 1, "name": "alpha", "routing_group": "slow", "provider": "anthropic", "kind": "quota_limited", "today_cost_usd": 0},
{"id": 2, "name": "beta", "routing_group": "sfast", "provider": "openai", "kind": "pay_as_you_go", "today_cost_usd": 0},
]
}
self.assertEqual([row["name"] for row in mod.normalize_account_rows(payload, "pay")], ["beta"])
self.assertEqual([row["name"] for row in mod.normalize_account_rows(payload, "alp")], ["alpha"])
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_infers_public_status_url_from_accounts_url(self) -> None:
mod = load_module()
self.assertEqual(
mod.inferred_status_url("https://codex.server2.shujk.top/1232131231313123/api/tui/accounts"),
"https://codex.server2.shujk.top/1232131231313123/api/status",
)
self.assertEqual(mod.inferred_status_url("https://example.com/nope"), "")
def test_version_update_message_only_for_newer_versions(self) -> None:
mod = load_module()
self.assertEqual(mod.latest_version_from_text('name = "shusub2"\nversion = "0.1.8"\n'), "0.1.8")
self.assertIn("0.1.7 -> 0.1.8", mod.version_update_message("0.1.8", "0.1.7"))
self.assertEqual(mod.version_update_message("0.1.7", "0.1.7"), "")
self.assertEqual(mod.version_update_message("0.1.6", "0.1.7"), "")
def test_write_api_url_config_uses_config_file_env(self) -> None:
mod = load_module()
with tempfile.TemporaryDirectory() as tmp:
config_file = Path(tmp) / "shusub2" / "api-url"
old_value = os.environ.get("SHUSUB2_API_URL_FILE")
os.environ["SHUSUB2_API_URL_FILE"] = str(config_file)
try:
written = mod.write_api_url_config("https://example.com/api/tui/accounts")
self.assertEqual(written, config_file)
self.assertEqual(config_file.read_text(encoding="utf-8"), "https://example.com/api/tui/accounts\n")
finally:
if old_value is None:
os.environ.pop("SHUSUB2_API_URL_FILE", None)
else:
os.environ["SHUSUB2_API_URL_FILE"] = old_value
def test_once_output_is_name_first_and_includes_daily_quota(self) -> None:
mod = load_module()
payload = {
"generated_at": "2026-06-09T12:00:00+08:00",
"totals": {"total_accounts": 1, "usable_accounts": 1, "today_cost_usd": 1, "today_tokens": 100, "today_requests": 2},
"accounts": [
{
"id": 2,
"name": "input 300",
"routing_group": "fast",
"provider": "openai",
"kind": "daily_limited",
"status": "ok",
"account_type": "apikey",
"daily_quota_limit": 300,
"daily_quota_used": 96.6,
"daily_quota_remaining": 203.4,
"daily_quota_used_percent": 32.2,
"daily_quota_reset_at": "2026-06-09T16:00:00Z",
"today_cost_usd": 1,
"today_tokens": 100,
"today_requests": 2,
"windows": [],
}
],
}
rows = mod.normalize_account_rows(payload)
self.assertEqual(rows[0]["kind_label"], "daily")
self.assertEqual(rows[0]["routing_group"], "fast")
self.assertEqual(rows[0]["provider"], "openai")
self.assertEqual(rows[0]["daily_quota_cell"], "96.6/300")
self.assertTrue(rows[0]["reset"].endswith("00:00") or rows[0]["reset"] != "-")
def test_daily_quota_cell_is_dash_without_limit(self) -> None:
mod = load_module()
self.assertEqual(mod.daily_quota_cell({"daily_quota_used": 0}), "-")
def test_monitor_availability_column_uses_bound_monitor_status(self) -> None:
mod = load_module()
payload = {
"generated_at": "2026-06-09T12:00:00+08:00",
"totals": {"total_accounts": 1, "usable_accounts": 1, "today_cost_usd": 1, "today_tokens": 100, "today_requests": 2},
"accounts": [
{
"id": 2,
"name": "input 300",
"base_url_hash": "input-hash",
"routing_group": "fast",
"provider": "openai",
"kind": "daily_limited",
"status": "ok",
"account_type": "apikey",
"today_cost_usd": 1,
"today_tokens": 100,
"today_requests": 2,
"windows": [],
}
],
}
status_payload = {
"channel_monitors": {
"items": [
{
"name": "other responses",
"base_url_hash": "input-hash",
"enabled": True,
"latest_status": "operational",
"latency_ms": 2606,
}
]
}
}
rows = mod.normalize_account_rows(payload)
out = io.StringIO()
with contextlib.redirect_stdout(out):
mod.print_once(payload, status_payload=status_payload)
self.assertEqual(mod.monitor_availability(rows[0], status_payload), "ok")
self.assertIn("availability", out.getvalue())
self.assertIn(" ok\n", out.getvalue())
def test_monitor_availability_is_dash_without_bound_monitor(self) -> None:
mod = load_module()
self.assertEqual(mod.monitor_availability({"name": "alpha"}, {}), "-")
def test_routing_group_sort_shows_sfast_first(self) -> None:
mod = load_module()
payload = {
"accounts": [
{"id": 1, "name": "manual", "routing_group": "id", "kind": "pay_as_you_go", "today_cost_usd": 10},
{"id": 2, "name": "slow", "routing_group": "slow", "kind": "pay_as_you_go", "today_cost_usd": 1},
{"id": 3, "name": "fast", "routing_group": "fast", "kind": "pay_as_you_go", "today_cost_usd": 1},
{"id": 4, "name": "sfast", "routing_group": "sfast", "kind": "pay_as_you_go", "today_cost_usd": 1},
]
}
rows = mod.normalize_account_rows(payload)
self.assertEqual([row["name"] for row in rows], ["sfast", "fast", "slow", "manual"])
def test_monitor_summary_and_account_match(self) -> None:
mod = load_module()
status_payload = {
"channel_monitors": {
"enabled": 6,
"latest_ok": 5,
"latest_failed": 0,
"latest_unknown": 1,
"latest_checked_max": "2026-06-09 15:07:40+08:00",
"items": [
{
"name": "input responses",
"provider": "openai",
"enabled": True,
"primary_model": "gpt-5.5",
"latest_status": "operational",
"latency_ms": 2606,
"checked_at": "2026-06-09 15:03:28+08:00",
"message": "",
},
{
"name": "kedaya responses",
"provider": "openai",
"enabled": True,
"primary_model": "gpt-5.5",
"latest_status": "degraded",
"latency_ms": 25923,
"checked_at": "2026-06-09 15:07:40+08:00",
"message": "slow response: 25923ms",
},
],
}
}
row = {"name": "input 300"}
self.assertIn("monitors 5 ok / 0 failed / 1 unknown", mod.monitor_summary(status_payload))
self.assertIn("monitor operational 2606ms", mod.monitor_detail(row, status_payload))
def test_monitor_match_prefers_base_url_hash_over_name_tokens(self) -> None:
mod = load_module()
status_payload = {
"channel_monitors": {
"items": [
{
"name": "input responses",
"base_url_hash": "wrong-hash",
"latest_status": "operational",
"latency_ms": 111,
},
{
"name": "unrelated responses",
"base_url_hash": "right-hash",
"latest_status": "degraded",
"latency_ms": 222,
},
],
}
}
row = {"name": "input 300", "base_url_hash": "right-hash"}
self.assertEqual(mod.matching_monitor(row, status_payload)["name"], "unrelated responses")
self.assertIn("monitor degraded 222ms", mod.monitor_detail(row, status_payload))
def test_monitor_error_is_non_blocking_in_once_output(self) -> None:
mod = load_module()
payload = {
"generated_at": "2026-06-09T12:00:00+08:00",
"totals": {"total_accounts": 1, "usable_accounts": 1, "today_cost_usd": 0, "today_tokens": 0, "today_requests": 0},
"accounts": [
{"id": 1, "name": "alpha", "routing_group": "slow", "provider": "openai", "kind": "pay_as_you_go", "status": "ok"}
],
}
out = io.StringIO()
with contextlib.redirect_stdout(out):
mod.print_once(payload, status_error="timeout")
text = out.getvalue()
self.assertIn("monitors error: timeout", text)
self.assertIn("alpha", text)
def test_once_output_includes_monitor_summary(self) -> None:
mod = load_module()
payload = {
"generated_at": "2026-06-09T12:00:00+08:00",
"totals": {"total_accounts": 0, "usable_accounts": 0, "today_cost_usd": 0, "today_tokens": 0, "today_requests": 0},
"accounts": [],
}
status_payload = {"channel_monitors": {"enabled": 1, "latest_ok": 1, "latest_failed": 0, "latest_unknown": 0}}
out = io.StringIO()
with contextlib.redirect_stdout(out):
mod.print_once(payload, status_payload=status_payload)
self.assertIn("monitors 1 ok / 0 failed / 0 unknown", out.getvalue())
def sample_logs_payload() -> dict:
return {
"code": 0,
"message": "success",
"data": {
"total": 2345,
"page": 1,
"page_size": 100,
"items": [
{
"id": 11,
"api_key_id": 3,
"api_key": {"id": 3, "name": "codex-main"},
"account_id": 7,
"account": {"id": 7, "name": "oai-sub-1"},
"user": {"id": 1, "name": "shujakuin"},
"model": "gpt-5.5",
"upstream_model": "gpt-5.5-codex",
"request_type": "stream",
"stream": True,
"input_tokens": 1200,
"output_tokens": 340,
"cache_creation_tokens": 50,
"cache_read_tokens": 4100,
"total_cost": 0.0123,
"actual_cost": 0.0061,
"duration_ms": 5321,
"first_token_ms": 800,
"request_id": "req-abc",
"created_at": "2026-07-21T10:00:00+08:00",
},
{
"id": 12,
"api_key_id": 4,
"account_id": 9,
"account": {"id": 9, "name": "claude-max"},
"model": "claude-sonnet-5",
"request_type": "",
"stream": False,
"input_tokens": 100,
"output_tokens": 20,
"total_cost": 0.001,
"duration_ms": 900,
"created_at": "2026-07-21T10:05:00+08:00",
},
],
},
}
class Sub2APILogsTests(unittest.TestCase):
def test_normalize_log_rows_maps_columns_and_sorts_newest_first(self) -> None:
mod = load_module()
rows = mod.normalize_log_rows(sample_logs_payload())
self.assertEqual([row["id"] for row in rows], [12, 11])
newest, oldest = rows
self.assertEqual(newest["key"], "#4")
self.assertEqual(newest["account"], "claude-max")
self.assertEqual(newest["type"], "sync")
self.assertEqual(newest["tokens"], 120)
self.assertEqual(oldest["key"], "codex-main")
self.assertEqual(oldest["account"], "oai-sub-1")
self.assertEqual(oldest["user"], "shujakuin")
self.assertEqual(oldest["type"], "stream")
self.assertEqual(oldest["tokens"], 1200 + 340 + 50 + 4100)
self.assertEqual(oldest["cost"], 0.0123)
self.assertEqual(oldest["duration_ms"], 5321)
self.assertTrue(oldest["time"].endswith("10:00") or oldest["time"] != "-")
def test_normalize_log_rows_filter_matches_key_account_model(self) -> None:
mod = load_module()
payload = sample_logs_payload()
self.assertEqual([row["id"] for row in mod.normalize_log_rows(payload, "codex-main")], [11])
self.assertEqual([row["id"] for row in mod.normalize_log_rows(payload, "claude")], [12])
self.assertEqual([row["id"] for row in mod.normalize_log_rows(payload, "stream")], [11])
self.assertEqual(mod.normalize_log_rows(payload, "no-match"), [])
def test_log_items_supports_flat_and_enveloped_payloads(self) -> None:
mod = load_module()
self.assertEqual(len(mod.log_items(sample_logs_payload())), 2)
self.assertEqual(len(mod.log_items({"items": [{"id": 1}]})), 1)
self.assertEqual(mod.log_items({"data": {"items": "nope"}}), [])
self.assertEqual(mod.log_items({}), [])
def test_log_type_label_falls_back_to_legacy_flags(self) -> None:
mod = load_module()
self.assertEqual(mod.log_type_label({"request_type": "ws_v2"}), "ws_v2")
self.assertEqual(mod.log_type_label({"request_type": "unknown", "openai_ws_mode": True}), "ws_v2")
self.assertEqual(mod.log_type_label({"stream": True}), "stream")
self.assertEqual(mod.log_type_label({"stream": False}), "sync")
self.assertEqual(mod.log_type_label({}), "-")
def test_logs_request_url_forces_first_page_and_limit(self) -> None:
mod = load_module()
url = mod.logs_request_url("https://sub2apicn.shujk.top/api/v1/admin/usage?page=9&page_size=5&user_id=2", 100)
parsed = dict(part.split("=") for part in url.split("?", 1)[1].split("&"))
self.assertEqual(parsed["page"], "1")
self.assertEqual(parsed["page_size"], "100")
self.assertEqual(parsed["user_id"], "2")
capped = mod.logs_request_url("https://sub2apicn.shujk.top/api/v1/admin/usage", 9000)
self.assertIn("page_size=1000", capped)
def test_log_tokens_per_second_uses_decode_window(self) -> None:
mod = load_module()
item = {"output_tokens": 340, "duration_ms": 5321, "first_token_ms": 800}
self.assertAlmostEqual(mod.log_tokens_per_second(item), 340 * 1000.0 / (5321 - 800))
self.assertEqual(mod.format_rate(mod.log_tokens_per_second(item)), "75.2/s")
# missing first token: falls back to the full duration
self.assertAlmostEqual(mod.log_tokens_per_second({"output_tokens": 100, "duration_ms": 2000}), 50.0)
# no output, zero duration, or first token >= duration -> no rate
self.assertEqual(mod.log_tokens_per_second({"output_tokens": 0, "duration_ms": 2000}), 0.0)
self.assertEqual(mod.log_tokens_per_second({"output_tokens": 10, "duration_ms": 0}), 0.0)
self.assertEqual(mod.log_tokens_per_second({"output_tokens": 10, "duration_ms": 500, "first_token_ms": 500}), 0.0)
self.assertEqual(mod.format_rate(0), "-")
self.assertEqual(mod.format_rate(123.4), "123/s")
def test_normalize_log_rows_carries_first_token_and_rate(self) -> None:
mod = load_module()
rows = mod.normalize_log_rows(sample_logs_payload())
oldest = rows[-1]
self.assertEqual(oldest["first_token_ms"], 800)
self.assertAlmostEqual(oldest["tokens_per_second"], 340 * 1000.0 / 4521)
self.assertIn("first 800ms", mod.log_detail_line(oldest))
self.assertIn("75.2/s", mod.log_detail_line(oldest))
def test_key_color_is_stable_and_from_palette(self) -> None:
mod = load_module()
first = mod.key_color("codex-main")
self.assertIn(first, mod.KEY_COLOR_PALETTE)
self.assertEqual(first, mod.key_color("codex-main"))
self.assertEqual(mod.key_color(""), "")
self.assertEqual(mod.key_color("-"), "")
colors = {mod.key_color(name) for name in ("a", "b", "c", "d", "e", "f")}
self.assertGreater(len(colors), 1)
def test_print_logs_once_renders_requested_columns(self) -> None:
mod = load_module()
out = io.StringIO()
with contextlib.redirect_stdout(out):
mod.print_logs_once(sample_logs_payload())
text = out.getvalue()
self.assertIn("key", text)
self.assertIn("codex-main", text)
self.assertIn("oai-sub-1", text)
self.assertIn("gpt-5.5", text)
self.assertIn("stream", text)
self.assertIn("5.7K", text)
self.assertIn("$0.012", text)
self.assertIn("5321ms", text)
self.assertIn("first", text)
self.assertIn("tok/s", text)
self.assertIn("800ms", text)
self.assertIn("75.2/s", text)
self.assertIn("total 2.3K records", text)
def test_default_logs_token_reads_config_file(self) -> None:
mod = load_module()
with tempfile.TemporaryDirectory() as tmp:
token_file = Path(tmp) / "logs-token"
token_file.write_text("# comment\nsk-admin-token\n", encoding="utf-8")
old_value = os.environ.get("SHUSUB2_LOGS_TOKEN_FILE")
os.environ["SHUSUB2_LOGS_TOKEN_FILE"] = str(token_file)
try:
self.assertEqual(mod.default_logs_token(), "sk-admin-token")
written = mod.write_config_value(mod.logs_token_config_file_path(), "sk-new-token")
self.assertEqual(written, token_file)
self.assertEqual(token_file.read_text(encoding="utf-8"), "sk-new-token\n")
finally:
if old_value is None:
os.environ.pop("SHUSUB2_LOGS_TOKEN_FILE", None)
else:
os.environ["SHUSUB2_LOGS_TOKEN_FILE"] = old_value
def test_once_logs_without_token_prints_hint_and_exits_2(self) -> None:
mod = load_module()
old_token = os.environ.pop("SHUSUB2_LOGS_TOKEN", None)
old_file = os.environ.get("SHUSUB2_LOGS_TOKEN_FILE")
os.environ["SHUSUB2_LOGS_TOKEN_FILE"] = "/nonexistent/shusub2/logs-token"
err = io.StringIO()
try:
with contextlib.redirect_stderr(err):
rc = mod.main(["--once", "--logs", "--no-version-check"])
finally:
if old_token is not None:
os.environ["SHUSUB2_LOGS_TOKEN"] = old_token
if old_file is None:
os.environ.pop("SHUSUB2_LOGS_TOKEN_FILE", None)
else:
os.environ["SHUSUB2_LOGS_TOKEN_FILE"] = old_file
self.assertEqual(rc, 2)
self.assertIn("logs token not configured", err.getvalue())
def test_default_logs_url_defaults_to_sub2apicn(self) -> None:
mod = load_module()
old_url = os.environ.pop("SHUSUB2_LOGS_URL", None)
old_file = os.environ.get("SHUSUB2_LOGS_URL_FILE")
os.environ["SHUSUB2_LOGS_URL_FILE"] = "/nonexistent/shusub2/logs-url"
try:
self.assertEqual(mod.default_logs_url(), "https://sub2apicn.shujk.top/api/v1/admin/usage")
finally:
if old_url is not None:
os.environ["SHUSUB2_LOGS_URL"] = old_url
if old_file is None:
os.environ.pop("SHUSUB2_LOGS_URL_FILE", None)
else:
os.environ["SHUSUB2_LOGS_URL_FILE"] = old_file
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink