server: 任务列表一次聚合回 total_tokens
_list_tasks 在取附件的同时用 _task_tokens_map 一次查 cost_records 按 task_id 求和(GROUP BY),一屏只多一条 查询,不按任务逐个调成本接口。无记录计 0、表不存在时 全计 0 且列表不失败;_task_to_dict 新增 total_tokens 参数,默认 None(非列表调用方不把“未算”写成 0)。 前端 TaskListSection 的 token 栏位此前已接好 total_tokens, 补上后列表行直接显示真数,不再显示 token —;前端注释 与占位文案同步更新。 新测试 test_task_list_tokens.py 3 项:HTTP 列表聚合 (actual+estimated 加总、无记录为 0)、缺表不炸、 _task_to_dict 默认 None。
This commit is contained in:
1 parent
65a875ada9
commit
9eedee5a3d
4 files changed
+168
-14
No files matched your search
@@ -0,0 +1,122 @@
|
||||
"""任务列表 total_tokens:一次聚合,不逐个调成本接口。
|
||||
|
||||
用真实 HTTP 服务 + 临时 SQLite:建两个任务,给其中一个记
|
||||
两条成本记录(actual + estimated),列表应一次回两行的
|
||||
total_tokens(有记录的加总、无记录的 0),且不回 404/500;
|
||||
_task_tokens_map 在 cost_records 表不存在时也不炸列表。
|
||||
"""
|
||||
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.cost import CostStore
|
||||
from cloud.server import (
|
||||
_task_to_dict, _task_tokens_map, create_server)
|
||||
except ImportError:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.cost import CostStore
|
||||
from cloud.server import (
|
||||
_task_to_dict, _task_tokens_map, create_server)
|
||||
|
||||
|
||||
class ListTokensTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
self.auth = AuthStore(self.db_path)
|
||||
self.owner_id, self.owner_key = self.auth.create_user("owner")
|
||||
self.server = create_server(self.db_path, "127.0.0.1", 0)
|
||||
self.port = self.server.server_address[1]
|
||||
self.thread = threading.Thread(
|
||||
target=self.server.serve_forever, daemon=True)
|
||||
self.thread.start()
|
||||
status, self.project = self._req(
|
||||
"POST", "/projects", {"name": "demo"}, self.owner_key)
|
||||
self.assertEqual(status, 201)
|
||||
|
||||
def tearDown(self):
|
||||
self.server.shutdown()
|
||||
self.server.server_close()
|
||||
self.thread.join(timeout=5)
|
||||
try:
|
||||
self.auth.conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _req(self, method, path, body=None, key=None):
|
||||
conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)
|
||||
hdrs = {}
|
||||
payload = None
|
||||
if body is not None:
|
||||
payload = json.dumps(body).encode("utf-8")
|
||||
hdrs["Content-Type"] = "application/json"
|
||||
if key is not None:
|
||||
hdrs["X-API-Key"] = key
|
||||
try:
|
||||
conn.request(method, path, body=payload, headers=hdrs)
|
||||
resp = conn.getresponse()
|
||||
data = resp.read()
|
||||
finally:
|
||||
conn.close()
|
||||
try:
|
||||
parsed = json.loads(data.decode("utf-8")) if data else None
|
||||
except ValueError:
|
||||
parsed = None
|
||||
return resp.status, parsed
|
||||
|
||||
def test_list_returns_aggregated_total_tokens(self):
|
||||
pid = self.project["id"]
|
||||
_, t1 = self._req("POST", "/tasks",
|
||||
{"project_id": pid, "prompt": "有成本"},
|
||||
self.owner_key)
|
||||
_, t2 = self._req("POST", "/tasks",
|
||||
{"project_id": pid, "prompt": "无成本"},
|
||||
self.owner_key)
|
||||
store = CostStore(self.db_path)
|
||||
try:
|
||||
store.record(t1["id"], pid, "m", input_tokens=100,
|
||||
output_tokens=50, source="actual")
|
||||
store.record(t1["id"], pid, "m", input_tokens=20,
|
||||
output_tokens=0, source="estimated")
|
||||
finally:
|
||||
store._conn.close()
|
||||
status, body = self._req(
|
||||
"GET", "/tasks?project_id=%s" % pid, key=self.owner_key)
|
||||
self.assertEqual(status, 200, body)
|
||||
by_id = {t["id"]: t for t in body["tasks"]}
|
||||
self.assertEqual(by_id[t1["id"]]["total_tokens"], 170)
|
||||
self.assertEqual(by_id[t2["id"]]["total_tokens"], 0)
|
||||
|
||||
def test_tokens_map_missing_table_is_zero_not_error(self):
|
||||
import sqlite3
|
||||
conn = sqlite3.connect(":memory:")
|
||||
try:
|
||||
self.assertEqual(
|
||||
_task_tokens_map(conn, ["t1", "t2"]),
|
||||
{"t1": 0, "t2": 0})
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def test_task_to_dict_default_total_tokens_is_none(self):
|
||||
row = {"id": "t1", "project_id": "p1", "prompt": "p",
|
||||
"status": "completed", "idempotency_key": None,
|
||||
"depends_on": "[]", "priority": 0, "created_by": "u",
|
||||
"created_at": 1.0, "updated_at": 2.0,
|
||||
"last_heartbeat": None}
|
||||
self.assertIsNone(_task_to_dict(row, [])["total_tokens"])
|
||||
self.assertEqual(
|
||||
_task_to_dict(row, [], 42)["total_tokens"], 42)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user