270 lines
11 KiB
Python
270 lines
11 KiB
Python
"""成本/配额/文件标记单测：task-cost / quota-status / task-files
|
|
解析、sandbox 提取、worker 投递与 worktree 解析。
|
|
|
|
不跑真实模型、不联网：解析是纯函数；worker 投递用 init_db
|
|
建的临时 SQLite 库，直接调 _deliver_marker_tools，不走
|
|
run_one 全量流程；文件读取用临时目录冒充 worktree（测
|
|
试里直接给 Worker 换 tenants_root 与 host_py_path）。
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
try:
|
|
from cloud.cost import CostStore
|
|
from cloud.marker_tools import (
|
|
extract_quota_status,
|
|
extract_task_costs,
|
|
extract_task_files,
|
|
format_cost,
|
|
safe_worktree_path,
|
|
)
|
|
from cloud.sandbox import LocalExecutor, _parse_host_result
|
|
from cloud.server import init_db
|
|
from cloud.worker import Worker
|
|
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.cost import CostStore
|
|
from cloud.marker_tools import (
|
|
extract_quota_status,
|
|
extract_task_costs,
|
|
extract_task_files,
|
|
format_cost,
|
|
safe_worktree_path,
|
|
)
|
|
from cloud.sandbox import LocalExecutor, _parse_host_result
|
|
from cloud.server import init_db
|
|
from cloud.worker import Worker
|
|
|
|
FAKE_ENV = {"XCM_API_KEY": "test-placeholder"}
|
|
|
|
|
|
class ExtractTestCase(unittest.TestCase):
|
|
def test_cost_and_files(self):
|
|
self.assertEqual(extract_task_costs(
|
|
"```task-cost\n{\"task_id\": \"t1\"}\n```"),
|
|
[{"task_id": "t1"}])
|
|
self.assertEqual(extract_task_costs(
|
|
"```task-cost\n{\"to\": \"t1\"}\n```"),
|
|
[{"task_id": "t1"}])
|
|
self.assertEqual(extract_task_costs(
|
|
"```task-cost\n{}\n```"), [])
|
|
self.assertEqual(extract_task_files(
|
|
"```task-files\n{\"task_id\": \"t1\"}\n```"),
|
|
[{"task_id": "t1", "path": ""}])
|
|
self.assertEqual(extract_task_files(
|
|
"```task-files\n{\"task_id\": \"t1\", \"path\": \"a.txt\"}\n```"),
|
|
[{"task_id": "t1", "path": "a.txt"}])
|
|
self.assertEqual(extract_task_files(
|
|
"```task-files\n{\"path\": \"a.txt\"}\n```"), [])
|
|
|
|
def test_quota_only_object(self):
|
|
self.assertEqual(extract_quota_status(
|
|
"```quota-status\n{}\n```"), [{}])
|
|
self.assertEqual(extract_quota_status(
|
|
"```quota-status\n{\"owner\": \"someone\"}\n```"), [{}])
|
|
self.assertEqual(extract_quota_status(
|
|
"```quota-status\n[]\n```"), [])
|
|
|
|
def test_safe_worktree_path(self):
|
|
self.assertIsNone(safe_worktree_path("/tmp/w", "../etc/passwd"))
|
|
self.assertIsNone(safe_worktree_path("/tmp/w", "/etc/passwd"))
|
|
self.assertEqual(safe_worktree_path("/tmp/w", ""),
|
|
"/tmp/w")
|
|
self.assertEqual(safe_worktree_path("/tmp/w", "a/b.txt"),
|
|
"/tmp/w/a/b.txt")
|
|
|
|
def test_format_cost_unpriced_separate(self):
|
|
line = format_cost("t1", {"usd_micros": 0, "total_tokens": 10,
|
|
"actual_tokens": 10,
|
|
"estimated_tokens": 0, "turns": 1,
|
|
"unpriced_actual_tokens": 10})
|
|
self.assertIn("$0.0000", line)
|
|
self.assertIn("无价 actual token 10 未计入美元", line)
|
|
|
|
def test_sandbox_extracts_both_executors(self):
|
|
text = ("```task-cost\n{\"task_id\": \"t1\"}\n```\n"
|
|
"```quota-status\n{}\n```\n"
|
|
"```task-files\n{\"task_id\": \"t1\"}\n```")
|
|
fake = SimpleNamespace(
|
|
returncode=0,
|
|
stdout=json.dumps({"status": "completed", "text": text}),
|
|
stderr="")
|
|
local = LocalExecutor("/fake/host.py")._result_from_runner("t1", fake)
|
|
self.assertEqual(len(local["task_costs"]), 1)
|
|
self.assertEqual(len(local["quota_status"]), 1)
|
|
self.assertEqual(len(local["task_files"]), 1)
|
|
docker = _parse_host_result("t1", fake)
|
|
self.assertEqual(len(docker["task_costs"]), 1)
|
|
self.assertEqual(len(docker["quota_status"]), 1)
|
|
self.assertEqual(len(docker["task_files"]), 1)
|
|
|
|
|
|
class WorkerCostFilesTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.db_path = os.path.join(self.tmp.name, "control.db")
|
|
init_db(self.db_path)
|
|
self._insert("PM", status="running", created_by="tester")
|
|
self._member("tester", "maintainer")
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def _insert(self, task_id, status="queued", created_by="tester",
|
|
project_id="p1", prompt="do work", parent=None,
|
|
**cols):
|
|
now = time.time()
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
"INSERT INTO tasks(id, project_id, prompt, status,"
|
|
" idempotency_key, depends_on, priority, created_by,"
|
|
" created_at, updated_at, last_heartbeat)"
|
|
" VALUES(?, ?, ?, ?, NULL, '[]', 0, ?, ?, ?, NULL)",
|
|
(task_id, project_id, prompt, status, created_by,
|
|
now, now))
|
|
if parent is not None:
|
|
conn.execute(
|
|
"UPDATE tasks SET parent_task_id = ? WHERE id = ?",
|
|
(parent, task_id))
|
|
for key, value in cols.items():
|
|
conn.execute(
|
|
"UPDATE tasks SET %s = ? WHERE id = ?" % key,
|
|
(value, task_id))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def _member(self, user_id, role):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
"CREATE TABLE IF NOT EXISTS users(id TEXT PRIMARY KEY,"
|
|
" username TEXT UNIQUE, api_key_hash TEXT,"
|
|
" created_at REAL)")
|
|
conn.execute(
|
|
"CREATE TABLE IF NOT EXISTS project_members("
|
|
"project_id TEXT, user_id TEXT, role TEXT,"
|
|
" PRIMARY KEY(project_id, user_id))")
|
|
conn.execute(
|
|
"INSERT OR REPLACE INTO project_members"
|
|
"(project_id, user_id, role) VALUES('p1', ?, ?)",
|
|
(user_id, role))
|
|
conn.execute(
|
|
"INSERT OR IGNORE INTO users(id, username, api_key_hash,"
|
|
" created_at) VALUES(?, ?, '', 0)", (user_id, user_id))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def _worker(self):
|
|
return Worker(self.db_path, "/fake/host.py", env=FAKE_ENV)
|
|
|
|
def _followup_bodies(self, task_id="PM"):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
rows = conn.execute(
|
|
"SELECT body FROM task_followups"
|
|
" WHERE task_id = ?", (task_id,)).fetchall()
|
|
return [row[0] for row in rows]
|
|
finally:
|
|
conn.close()
|
|
|
|
def test_deliver_cost_and_quota(self):
|
|
self._insert("E1", status="completed", parent="PM")
|
|
store = CostStore(self.db_path)
|
|
try:
|
|
store.record("E1", "p1", "gpt-test", 100, 50,
|
|
usd_micros=12345, source="actual")
|
|
finally:
|
|
try:
|
|
store._conn.close()
|
|
except Exception:
|
|
pass
|
|
worker = self._worker()
|
|
worker._deliver_marker_tools(worker.get_task("PM"), {
|
|
"status": "completed",
|
|
"task_costs": [{"task_id": "E1"}],
|
|
"quota_status": [{}],
|
|
"task_files": []})
|
|
bodies = self._followup_bodies()
|
|
self.assertEqual(len(bodies), 1)
|
|
self.assertIn("$0.0123", bodies[0])
|
|
self.assertIn("配额", bodies[0])
|
|
|
|
def test_deliver_cost_other_creator_rejected(self):
|
|
self._insert("E1", status="completed", created_by="someone-else")
|
|
worker = self._worker()
|
|
worker._deliver_marker_tools(worker.get_task("PM"), {
|
|
"status": "completed",
|
|
"task_costs": [{"task_id": "E1"}],
|
|
"quota_status": [],
|
|
"task_files": []})
|
|
bodies = self._followup_bodies()
|
|
self.assertIn("不在派生链内且创建人不同", bodies[0])
|
|
|
|
def test_deliver_files_list_and_read(self):
|
|
self._insert("E1", status="completed", parent="PM")
|
|
worker = self._worker()
|
|
# 冒充 worktree：租户路径下给 E1 建目录与文件。
|
|
worktree = os.path.join(self.tmp.name, "tenants", "p1",
|
|
"work", "host-worktrees", "E1")
|
|
os.makedirs(worktree)
|
|
with open(os.path.join(worktree, "out.txt"), "w") as fh:
|
|
fh.write("hello-files")
|
|
with open(os.path.join(worktree, "bin.dat"), "wb") as fh:
|
|
fh.write(b"\x00\x01\x02")
|
|
worker.tenants_root = os.path.join(self.tmp.name, "tenants")
|
|
worker._deliver_marker_tools(worker.get_task("PM"), {
|
|
"status": "completed",
|
|
"task_costs": [],
|
|
"quota_status": [],
|
|
"task_files": [{"task_id": "E1", "path": ""},
|
|
{"task_id": "E1", "path": "out.txt"},
|
|
{"task_id": "E1", "path": "bin.dat"},
|
|
{"task_id": "E1", "path": "../escape"}]})
|
|
bodies = self._followup_bodies()
|
|
self.assertEqual(len(bodies), 1)
|
|
self.assertIn("out.txt", bodies[0])
|
|
self.assertIn("hello-files", bodies[0])
|
|
self.assertIn("二进制文件", bodies[0])
|
|
self.assertIn("路径逃逸", bodies[0])
|
|
|
|
def test_deliver_files_cloud_and_no_worktree_rejected(self):
|
|
self._insert("C1", status="completed", parent="PM",
|
|
executor="cloud")
|
|
self._insert("N1", status="completed", parent="PM")
|
|
worker = self._worker()
|
|
worker.tenants_root = os.path.join(self.tmp.name, "tenants")
|
|
worker._deliver_marker_tools(worker.get_task("PM"), {
|
|
"status": "completed",
|
|
"task_costs": [],
|
|
"quota_status": [],
|
|
"task_files": [{"task_id": "C1", "path": ""},
|
|
{"task_id": "N1", "path": ""}]})
|
|
bodies = self._followup_bodies()
|
|
self.assertIn("不在本机", bodies[0])
|
|
self.assertIn("没有本地 worktree", bodies[0])
|
|
|
|
def test_deliver_goal_sender_ignored(self):
|
|
self._insert("G1", status="running", goal_mode=1)
|
|
worker = self._worker()
|
|
worker._deliver_marker_tools(worker.get_task("G1"), {
|
|
"status": "completed",
|
|
"task_costs": [{"task_id": "G1"}],
|
|
"quota_status": [{}],
|
|
"task_files": []})
|
|
self.assertEqual(self._followup_bodies("G1"), [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|