382 lines
16 KiB
Python
382 lines
16 KiB
Python
"""覆盖率 78%->85% 第 4 步：cloud/worker.py _deliver_marker_tools
|
||
剩余分支的真实行为测试。
|
||
|
||
父任务属主用 AuthStore 建真实 owner + 项目成员关系（权限
|
||
True/False 两路都走真实 check_permission）；待办/自动化走
|
||
真实 SQLite 表；任务文件走临时 host.py 同级 work 目录。
|
||
绝不跑模型、不碰真实网关。
|
||
"""
|
||
|
||
import os
|
||
import sqlite3
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
|
||
from cloud.auth import AuthStore
|
||
from cloud.server import init_db
|
||
from cloud.worker import Worker
|
||
|
||
|
||
def _insert_task(db_path, task_id, status="queued", **kw):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
cols = {"id": task_id, "project_id": kw.pop("project_id", "proj"),
|
||
"prompt": kw.pop("prompt", "do it"), "status": status,
|
||
"created_by": kw.pop("created_by", "owner1"),
|
||
"created_at": kw.pop("created_at", 1000.0),
|
||
"updated_at": kw.pop("updated_at", 1000.0)}
|
||
cols.update(kw)
|
||
names = ", ".join(cols)
|
||
placeholders = ", ".join("?" * len(cols))
|
||
conn.execute("INSERT INTO tasks(%s) VALUES(%s)"
|
||
% (names, placeholders), tuple(cols.values()))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _rows(db_path, sql, params=()):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
return conn.execute(sql, params).fetchall()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _set(db_path, sql, params=()):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute(sql, params)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
class MarkerBase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.addCleanup(self.tmp.cleanup)
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
init_db(self.db_path)
|
||
auth = AuthStore(self.db_path)
|
||
self.owner_id, _ = auth.create_user("owner")
|
||
auth.create_project("proj", self.owner_id)
|
||
self.outsider_id, _ = auth.create_user("outsider")
|
||
auth.conn.close()
|
||
# host.py 替身放在临时目录：任务文件解析走它同级的
|
||
# work/host-worktrees/<任务名>，不碰真实 /fake 或仓库。
|
||
self.fake_host = os.path.join(self.tmp.name, "host.py")
|
||
with open(self.fake_host, "w") as f:
|
||
f.write("# fake host\n")
|
||
self.worker = Worker(
|
||
self.db_path, self.fake_host,
|
||
env={"XCM_API_KEY": "test-placeholder"},
|
||
heartbeat_file=os.path.join(self.tmp.name, "hb.json"),
|
||
worker_id="cov85-worker4", max_retries=0,
|
||
retry_backoff_seconds=0)
|
||
|
||
def _parent(self, pid="parent1", creator=None):
|
||
if self.worker.get_task(pid) is None:
|
||
_insert_task(self.db_path, pid, status="completed",
|
||
created_by=creator or self.owner_id)
|
||
return self.worker.get_task(pid)
|
||
|
||
def _followups(self, pid):
|
||
return [r[0] for r in _rows(
|
||
self.db_path,
|
||
"SELECT body FROM task_followups WHERE task_id = ?",
|
||
(pid,))]
|
||
|
||
def _deliver(self, pid, result, creator=None):
|
||
parent = self._parent(pid, creator)
|
||
self.worker._deliver_marker_tools(parent, result)
|
||
return self._followups(pid)
|
||
|
||
|
||
class CreatesGateTest(MarkerBase):
|
||
def test_creates_quota_reject(self):
|
||
_set(self.db_path,
|
||
"INSERT INTO user_quotas(user_id, quota_usd_micros,"
|
||
" updated_at) VALUES(?, 0, 0)", (self.owner_id,))
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_creates": [{"prompt": "x", "name": "c",
|
||
"model": None, "effort": None,
|
||
"executor": None}]})
|
||
self.assertEqual(len(bodies), 1)
|
||
self.assertIn("创建跳过", bodies[0])
|
||
self.assertEqual(_rows(
|
||
self.db_path,
|
||
"SELECT COUNT(*) FROM tasks WHERE parent_task_id = 'parent1'"),
|
||
[(0,)])
|
||
|
||
def test_creates_depth_and_direct_caps(self):
|
||
# 直接派生上限：_MARKER_MAX_DIRECT 个子任务已存在时
|
||
# 再创建被拒（报告写清上限，不炸）。
|
||
from cloud.worker import _MARKER_MAX_DIRECT
|
||
for i in range(_MARKER_MAX_DIRECT):
|
||
_insert_task(self.db_path, "kid-%d" % i, status="queued",
|
||
parent_task_id="parent1",
|
||
created_by=self.owner_id)
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_creates": [{"prompt": "x", "name": "c",
|
||
"model": None, "effort": None,
|
||
"executor": None}]})
|
||
self.assertIn("直接派生子任务已达上限", bodies[0])
|
||
|
||
|
||
class AutomationCreateTest(MarkerBase):
|
||
def test_automation_auth_reject(self):
|
||
# 外人（非成员）当父任务属主：权限 False，整轮跳过。
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"automation_creates": [{"prompt": "巡检", "name": "a1",
|
||
"interval_seconds": 3600}]},
|
||
creator=self.outsider_id)
|
||
self.assertIn("自动化创建跳过", bodies[0])
|
||
self.assertIn("无 create_task 权限", bodies[0])
|
||
|
||
def test_automation_past_first_run_reject_and_success(self):
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"automation_creates": [
|
||
{"prompt": "旧的", "name": "old",
|
||
"interval_seconds": 3600,
|
||
"first_run_at": time.time() - 3600},
|
||
{"prompt": "巡检任务", "name": "new",
|
||
"interval_seconds": 3600,
|
||
"first_run_at": time.time() + 3600}]})
|
||
self.assertIn("首次运行时刻已过去", bodies[0])
|
||
self.assertIn("已创建自动化", bodies[0])
|
||
self.assertEqual(_rows(
|
||
self.db_path,
|
||
"SELECT COUNT(*) FROM automations WHERE project_id = 'proj'"),
|
||
[(1,)])
|
||
|
||
def test_automation_quota_reject(self):
|
||
_set(self.db_path,
|
||
"INSERT INTO user_quotas(user_id, quota_usd_micros,"
|
||
" updated_at) VALUES(?, 0, 0)", (self.owner_id,))
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"automation_creates": [{"prompt": "巡检", "name": "a1",
|
||
"interval_seconds": 3600}]})
|
||
self.assertIn("自动化创建跳过", bodies[0])
|
||
self.assertIn("配额检查未通过", bodies[0])
|
||
|
||
|
||
class CancelRetryTest(MarkerBase):
|
||
def test_cancel_descendant_and_retry_failed(self):
|
||
_insert_task(self.db_path, "kid1", status="queued",
|
||
parent_task_id="parent1",
|
||
created_by=self.owner_id)
|
||
_insert_task(self.db_path, "kid2", status="failed",
|
||
parent_task_id="parent1",
|
||
created_by=self.owner_id)
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_cancels": [{"task_id": "kid1"}],
|
||
"task_retries": [{"task_id": "kid2"}]})
|
||
self.assertIn("已取消任务 kid1", bodies[0])
|
||
self.assertIn("已重试任务 kid2", bodies[0])
|
||
statuses = dict(_rows(self.db_path,
|
||
"SELECT id, status FROM tasks"))
|
||
self.assertEqual(statuses["kid1"], "interrupted")
|
||
self.assertEqual(statuses["kid2"], "queued")
|
||
|
||
def test_cancel_other_project_rejected(self):
|
||
_insert_task(self.db_path, "alien", status="queued",
|
||
project_id="other", created_by="someone-else")
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_cancels": [{"task_id": "alien"}]})
|
||
self.assertIn("拒绝", bodies[0])
|
||
self.assertEqual(_rows(
|
||
self.db_path,
|
||
"SELECT status FROM tasks WHERE id = 'alien'"),
|
||
[("queued",)])
|
||
|
||
def test_cancel_missing_target_rejected(self):
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_cancels": [{"task_id": "ghost"}]})
|
||
self.assertIn("拒绝", bodies[0])
|
||
|
||
|
||
class TodoTaskBlocksTest(MarkerBase):
|
||
def test_todo_create_update_list(self):
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"todo_creates": [{"title": "写周报", "description": "周五前",
|
||
"status": "pending",
|
||
"priority": "high"}]})
|
||
self.assertIn("已创建待办", bodies[0])
|
||
todo_id = _rows(self.db_path,
|
||
"SELECT id FROM todos")[0][0]
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"todo_updates": [{"id": todo_id,
|
||
"fields": {"status": "in_progress"}}],
|
||
"todo_lists": [{"status": "in_progress"}]})
|
||
# 第二次 deliver 追加第二条报告：更新与列表都在 bodies[1]。
|
||
self.assertIn("已更新待办", bodies[1])
|
||
self.assertIn("待办列表", bodies[1])
|
||
self.assertIn("写周报", bodies[1])
|
||
|
||
def test_todo_auth_reject_and_bad_refs(self):
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"todo_creates": [{"title": "挂别人任务",
|
||
"task_id": "ghost"}],
|
||
"todo_updates": [{"id": "ghost-todo",
|
||
"fields": {"status": "done"}}],
|
||
"todo_lists": [{}]}, creator=self.outsider_id)
|
||
self.assertIn("拒绝", bodies[0])
|
||
# 外人无 read/create：三类块都被拒，待办表根本未建。
|
||
self.assertEqual(_rows(
|
||
self.db_path,
|
||
"SELECT name FROM sqlite_master WHERE type = 'table'"
|
||
" AND name = 'todos'"), [])
|
||
|
||
def test_todo_create_cross_project_ref_rejected(self):
|
||
_insert_task(self.db_path, "alien", status="queued",
|
||
project_id="other", created_by="someone-else")
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"todo_creates": [{"title": "挂跨项目",
|
||
"task_id": "alien"}]})
|
||
self.assertIn("拒绝", bodies[0])
|
||
|
||
def test_task_list_filters_and_cost_and_quota(self):
|
||
_insert_task(self.db_path, "kid1", status="completed",
|
||
parent_task_id="parent1",
|
||
created_by=self.owner_id, prompt="做加法")
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_lists": [{"status": "completed", "q": "加法",
|
||
"limit": 10}],
|
||
"task_costs": [{"task_id": "kid1"}],
|
||
"quota_status": [{}]})
|
||
self.assertIn("任务列表", bodies[0])
|
||
self.assertIn("kid1", bodies[0])
|
||
self.assertIn("配额", bodies[0])
|
||
|
||
def test_task_list_empty_and_cost_missing(self):
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_lists": [{"status": "failed", "q": None,
|
||
"limit": 10}],
|
||
"task_costs": [{"task_id": "ghost"}]})
|
||
self.assertIn("任务列表：无匹配项", bodies[0])
|
||
self.assertIn("不存在", bodies[0])
|
||
|
||
def test_quota_status_no_creator_rejected(self):
|
||
# created_by 为空：_marker_auth 返回 False，read 路先被拒。
|
||
_insert_task(self.db_path, "parent1", status="completed",
|
||
created_by="")
|
||
parent = self.worker.get_task("parent1")
|
||
self.worker._deliver_marker_tools(parent, {
|
||
"status": "completed", "quota_status": [{}]})
|
||
bodies = self._followups("parent1")
|
||
self.assertIn("拒绝", bodies[0])
|
||
|
||
|
||
class TaskFilesTest(MarkerBase):
|
||
def _worktree(self, tid):
|
||
d = os.path.join(self.tmp.name, "work", "host-worktrees", tid)
|
||
os.makedirs(d, exist_ok=True)
|
||
return d
|
||
|
||
def test_file_read_dir_list_and_missing(self):
|
||
d = self._worktree("kid1")
|
||
with open(os.path.join(d, "a.txt"), "w") as f:
|
||
f.write("hello file")
|
||
with open(os.path.join(d, "b.bin"), "wb") as f:
|
||
f.write(b"\x00\x01\x02binary")
|
||
_insert_task(self.db_path, "kid1", status="completed",
|
||
parent_task_id="parent1",
|
||
created_by=self.owner_id)
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_files": [{"task_id": "kid1", "path": "a.txt"},
|
||
{"task_id": "kid1", "path": ""},
|
||
{"task_id": "kid1", "path": "nope.txt"},
|
||
{"task_id": "kid1", "path": "b.bin"},
|
||
{"task_id": "kid1", "path": "../evil"}]})
|
||
body = bodies[0]
|
||
self.assertIn("hello file", body)
|
||
self.assertIn("a.txt", body)
|
||
self.assertIn("不存在", body)
|
||
self.assertIn("二进制文件", body)
|
||
self.assertIn("路径逃逸", body)
|
||
|
||
def test_file_no_worktree_rejected(self):
|
||
_insert_task(self.db_path, "kid1", status="completed",
|
||
parent_task_id="parent1",
|
||
created_by=self.owner_id)
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_files": [{"task_id": "kid1", "path": "a.txt"}]})
|
||
self.assertIn("没有本地 worktree", bodies[0])
|
||
|
||
def test_file_auth_reject(self):
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_files": [{"task_id": "kid1", "path": "a.txt"}]},
|
||
creator=self.outsider_id)
|
||
self.assertIn("拒绝", bodies[0])
|
||
|
||
|
||
class PollWaitTest(MarkerBase):
|
||
def test_wait_mixed_targets(self):
|
||
_insert_task(self.db_path, "kid1", status="completed",
|
||
parent_task_id="parent1",
|
||
created_by=self.owner_id)
|
||
_insert_task(self.db_path, "alien", status="completed",
|
||
project_id="other", created_by="someone-else")
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"wait_threads": [{"task_ids": ["kid1", "ghost", "alien",
|
||
"parent1"],
|
||
"interval_seconds": 0.05,
|
||
"timeout_seconds": 5, "chars": 200}]})
|
||
body = bodies[0]
|
||
self.assertIn("kid1", body)
|
||
self.assertIn("不存在", body)
|
||
self.assertIn("跨项目", body)
|
||
self.assertIn("不能等待自己", body)
|
||
|
||
def test_poll_self_and_cross_project(self):
|
||
_insert_task(self.db_path, "alien", status="completed",
|
||
project_id="other", created_by="someone-else")
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"task_polls": [{"task_id": "parent1",
|
||
"interval_seconds": 0.05,
|
||
"timeout_seconds": 5, "chars": 200},
|
||
{"task_id": "alien",
|
||
"interval_seconds": 0.05,
|
||
"timeout_seconds": 5, "chars": 200},
|
||
{"task_id": "ghost",
|
||
"interval_seconds": 0.05,
|
||
"timeout_seconds": 5, "chars": 200}]})
|
||
body = bodies[0]
|
||
self.assertIn("不能轮询自己", body)
|
||
self.assertIn("跨项目", body)
|
||
self.assertIn("不存在", body)
|
||
|
||
def test_goal_convert_queues_goal_followup_after_report(self):
|
||
bodies = self._deliver("parent1", {
|
||
"status": "completed",
|
||
"todo_lists": [{}],
|
||
"goal_converts": [{"goal": "收齐周报"}]})
|
||
self.assertEqual(len(bodies), 2)
|
||
self.assertIn("待办列表", bodies[0])
|
||
self.assertEqual(bodies[1], "/goal 收齐周报")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|