"""core/storage/usage_report.py 的 DB 级测试 —— 锁「计费读侧口径」。 背景:units JSONB 的写侧(usage.py)与读侧(usage_report.py)是同一契约的两半, 此前读侧散在 web 层硬编码、零测试,改 key 会静默算错。本测试用真实 PG 验证: - cost 全 kind 合计、token/cache_hit 仅 chat - task 维度批量聚合 / 按模型聚合 / 按用户分页聚合 三条读路径口径一致 无 DB(本机 PG 没起 / 未配)则整组 skip,不拖累纯单测环境。 数据纪律(公测期):只 INSERT 自己造的 user/task/usage_events,teardown 只 DELETE 这些行(按测试专属 user_id 过滤),绝不触碰既有数据。 """ import os import unittest import uuid from datetime import datetime, timedelta, timezone from decimal import Decimal def _test_db_ready() -> bool: """DB 级测试只认显式 `ZCBOT_TEST_DB_URL`,绝不回退 .env 的 ZCBOT_DB_URL。 教训(2026-07-23 实锤):.env 的 ZCBOT_DB_URL 可能经隧道指向**生产库**(本次 即 127.0.0.1:6012 → 生产 PG),测试若沿用它,插入的 enabled 到点 job 会被生产 实例的调度守护认领并真跑 agent。测试库须显式指定: set ZCBOT_TEST_DB_URL=postgresql+psycopg://user:pass@127.0.0.1:5432/zcbot_test (建库后 `alembic upgrade head` 一次);未设 → 整组 skip。 """ url = os.environ.get("ZCBOT_TEST_DB_URL", "").strip() if not url: return False if "connect_timeout" not in url: # 测试库挂了要快速降级 skip,不能让整个 discovery 挂死在 TCP 建连上 url += ("&" if "?" in url else "?") + "connect_timeout=5" os.environ["ZCBOT_DB_URL"] = url # 本测试进程内覆盖,engine 单例随之指向测试库 return True try: if not _test_db_ready(): raise RuntimeError("ZCBOT_TEST_DB_URL 未设") from core.storage import session_scope from core.storage.models import Task, UsageEvent, User from core.storage import usage_report with session_scope() as _s: _s.execute(__import__("sqlalchemy").select(1)) _DB_OK = True except Exception: _DB_OK = False @unittest.skipUnless(_DB_OK, "ZCBOT_TEST_DB_URL 未设或测试库不可达,跳过 DB 级测试(绝不用 .env 的库)") class UsageReportTests(unittest.TestCase): """一个测试专属 user + 两个 task,插一组已知 usage_events,验证三条读路径。""" @classmethod def setUpClass(cls): cls.uid = uuid.uuid4() cls.tid_a = uuid.uuid4() cls.tid_b = uuid.uuid4() with session_scope() as s: s.add(User(user_id=cls.uid, email=f"test-usage-report-{cls.uid.hex[:8]}@invalid.local")) s.flush() # 无 relationship 映射,FK 依赖顺序要显式 flush 保证 for tid, name in ((cls.tid_a, "ur-test-a"), (cls.tid_b, "ur-test-b")): s.add(Task( task_id=tid, user_id=cls.uid, name=name, working_dir=f"workspace/users/{cls.uid}/{name}", )) s.flush() # task A:两笔 chat(带缓存命中)+ 一笔 image(cost-only,token 不该被计入) s.add(UsageEvent( user_id=cls.uid, task_id=cls.tid_a, kind="chat", model_profile="ur-test.flash", units={"tokens_in": 1000, "tokens_out": 200, "cache_hit_tokens": 600}, cost_cny=Decimal("0.010000"), )) s.add(UsageEvent( user_id=cls.uid, task_id=cls.tid_a, kind="chat", model_profile="ur-test.pro", units={"tokens_in": 3000, "tokens_out": 800, "cache_hit_tokens": 0}, cost_cny=Decimal("0.200000"), )) s.add(UsageEvent( user_id=cls.uid, task_id=cls.tid_a, kind="image", model_profile="ur-test-seedream", units={"images": 1, "tokens_in": 999999}, # 非 chat 的 tokens 字段必须被忽略 cost_cny=Decimal("0.300000"), )) # task B:一笔 chat;另一笔 task_id=NULL 的 kb_ingest(0022,不得进任何 task 聚合) s.add(UsageEvent( user_id=cls.uid, task_id=cls.tid_b, kind="chat", model_profile="ur-test.flash", units={"tokens_in": 500, "tokens_out": 100, "cache_hit_tokens": 250}, cost_cny=Decimal("0.005000"), )) s.add(UsageEvent( user_id=cls.uid, task_id=None, kind="kb_ingest", model_profile="ur-test.flash", units={"kb": "测试库", "source": "x.pdf", "tokens_in": 42, "tokens_out": 7}, cost_cny=Decimal("0.001000"), )) @classmethod def tearDownClass(cls): from sqlalchemy import delete with session_scope() as s: s.execute(delete(UsageEvent).where(UsageEvent.user_id == cls.uid)) s.execute(delete(Task).where(Task.user_id == cls.uid)) s.execute(delete(User).where(User.user_id == cls.uid)) def test_task_aggregates_cost_all_kinds_tokens_chat_only(self): with session_scope() as s: agg = usage_report.task_usage_aggregates(s, [self.tid_a, self.tid_b]) a = agg[self.tid_a] # cost = chat 0.01 + 0.20 + image 0.30;token 仅 chat(image 的 999999 忽略) self.assertAlmostEqual(a["cost_cny"], 0.51, places=6) self.assertEqual(a["tokens_in"], 4000) self.assertEqual(a["tokens_out"], 1000) self.assertEqual(a["tokens_cache_hit"], 600) b = agg[self.tid_b] self.assertAlmostEqual(b["cost_cny"], 0.005, places=6) self.assertEqual(b["tokens_in"], 500) self.assertEqual(b["tokens_cache_hit"], 250) def test_task_aggregates_empty_and_unknown(self): with session_scope() as s: self.assertEqual(usage_report.task_usage_aggregates(s, []), {}) # 未知 task 无行 → 不出现在结果里(调用方 .get 兜 0) self.assertEqual(usage_report.task_usage_aggregates(s, [uuid.uuid4()]), {}) def test_models_usage_groups_and_sorts(self): cutoff = datetime.now(timezone.utc) - timedelta(days=1) with session_scope() as s: rows = usage_report.models_usage(s, cutoff, sort="cost") by_mp = {r["model_profile"]: r for r in rows} # ur-test.flash:两笔 chat(task A 1000/200 + task B 500/100)+ 一笔 kb_ingest #(cost 计入、tokens 因非 chat 忽略) flash = by_mp["ur-test.flash"] self.assertEqual(flash["tokens_in"], 1500) self.assertEqual(flash["tokens_out"], 300) self.assertAlmostEqual(flash["cost_cny"], 0.016, places=6) self.assertEqual(flash["n_events"], 3) self.assertAlmostEqual(by_mp["ur-test.pro"]["cost_cny"], 0.2, places=6) self.assertAlmostEqual(by_mp["ur-test-seedream"]["cost_cny"], 0.3, places=6) self.assertEqual(by_mp["ur-test-seedream"]["tokens_in"], 0) def test_models_usage_cutoff_excludes_old(self): # cutoff 在未来 → 我们刚插的行全部被排除 future = datetime.now(timezone.utc) + timedelta(days=1) with session_scope() as s: rows = usage_report.models_usage(s, future, sort="cost") self.assertNotIn("ur-test.flash", {r["model_profile"] for r in rows}) def test_user_usage_page_finds_our_user(self): with session_scope() as s: d = usage_report.user_usage_page(s, page=0, page_size=100000, cutoff=None, sort="cost") mine = [r for r in d["rows"] if r["user_id"] == str(self.uid)] self.assertEqual(len(mine), 1) r = mine[0] # cost 全 kind:0.01+0.20+0.30+0.005+0.001;token 仅 chat:4500/1100;hit:850 self.assertAlmostEqual(r["cost_cny"], 0.516, places=6) self.assertEqual(r["tokens_in"], 4500) self.assertEqual(r["tokens_out"], 1100) self.assertEqual(r["tokens_cache_hit"], 850) self.assertEqual(r["n_events"], 5) self.assertIsNotNone(r["last_used_at"]) def test_usage_overview_delta(self): """overview 是全局聚合,共享 dev 库上断相对增量:排除我们行前后的差 = 我们插入的量。""" cutoff_7d = datetime.now(timezone.utc) - timedelta(days=7) with session_scope() as s: total = usage_report.usage_overview(s, cutoff_7d)["total"] # 我们贡献:cost 0.516 / tokens_in 4500 / tokens_out 1100 / hit 850 / 5 事件。 # 断"至少包含"(库里还有真实数据,只验下界与口径不炸)。 self.assertGreaterEqual(total["cost_cny"], 0.516 - 1e-6) self.assertGreaterEqual(total["tokens_in"], 4500) self.assertGreaterEqual(total["n_events"], 5) self.assertIsInstance(total["tokens_cache_hit"], int) if __name__ == "__main__": unittest.main()