180 lines
8.5 KiB
Python
180 lines
8.5 KiB
Python
"""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
|
|
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()
|