cloud: production control plane (API, auth/RBAC, queue, sandbox, memory/KB, fallback/cost, webhooks)
This commit is contained in:
1 parent
b97c34465f
commit
b74d99e959
17 files changed
+3420
No files matched your search
@@ -0,0 +1,149 @@
|
||||
"""auth 模块单测:API key 认证、角色权限、审计只增不减。"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from cloud.auth import ROLES, AuthStore
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.auth import ROLES, AuthStore
|
||||
|
||||
|
||||
class TestAuth(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.store = AuthStore(os.path.join(self.tmp.name, "auth.db"))
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _user(self, username):
|
||||
return self.store.create_user(username)
|
||||
|
||||
def _project_with_roles(self):
|
||||
"""owner + viewer + contributor + maintainer + 一个非成员,返回其 id。"""
|
||||
owner_id, _ = self._user("owner")
|
||||
viewer_id, _ = self._user("viewer")
|
||||
contributor_id, _ = self._user("contributor")
|
||||
maintainer_id, _ = self._user("maintainer")
|
||||
outsider_id, _ = self._user("outsider")
|
||||
project_id = "proj-1"
|
||||
self.store.create_project(project_id, owner_id)
|
||||
self.store.set_member(project_id, owner_id, viewer_id, "viewer")
|
||||
self.store.set_member(project_id, owner_id, contributor_id, "contributor")
|
||||
self.store.set_member(project_id, owner_id, maintainer_id, "maintainer")
|
||||
return project_id, owner_id, viewer_id, contributor_id, maintainer_id, outsider_id
|
||||
|
||||
def test_create_user_plain_key_authenticates(self):
|
||||
user_id, api_key_plain = self._user("alice")
|
||||
self.assertEqual(self.store.authenticate(api_key_plain), user_id)
|
||||
|
||||
def test_wrong_key_rejected(self):
|
||||
user_id, api_key_plain = self._user("alice")
|
||||
self.assertIsNone(self.store.authenticate(api_key_plain + "x"))
|
||||
self.assertIsNone(self.store.authenticate("not-a-real-key"))
|
||||
self.assertIsNone(self.store.authenticate(""))
|
||||
self.assertIsNone(self.store.authenticate(None))
|
||||
|
||||
def test_hash_stored_not_plaintext(self):
|
||||
_user_id, api_key_plain = self._user("alice")
|
||||
db_path = os.path.join(self.tmp.name, "auth.db")
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT api_key_hash FROM users WHERE username = ?", ("alice",)
|
||||
).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertIsNotNone(row)
|
||||
stored = row[0]
|
||||
self.assertNotEqual(stored, api_key_plain)
|
||||
self.assertNotIn(api_key_plain, stored)
|
||||
# 库里只有哈希列,没有明文列可查
|
||||
columns = [c[1] for c in sqlite3.connect(db_path).execute(
|
||||
"PRAGMA table_info(users)").fetchall()]
|
||||
self.assertNotIn("api_key_plain", columns)
|
||||
self.assertNotIn("api_key", columns)
|
||||
|
||||
def test_permission_denials(self):
|
||||
project_id, _o, viewer_id, contributor_id, _m, outsider_id = (
|
||||
self._project_with_roles()
|
||||
)
|
||||
# viewer 不能 create_task
|
||||
self.assertFalse(
|
||||
self.store.check_permission(viewer_id, project_id, "create_task"))
|
||||
with self.assertRaises(PermissionError):
|
||||
self.store.require_permission(viewer_id, project_id, "create_task")
|
||||
# contributor 不能 approve_merge
|
||||
self.assertFalse(
|
||||
self.store.check_permission(contributor_id, project_id, "approve_merge"))
|
||||
with self.assertRaises(PermissionError):
|
||||
self.store.require_permission(
|
||||
contributor_id, project_id, "approve_merge")
|
||||
# 非成员读不到
|
||||
self.assertFalse(
|
||||
self.store.check_permission(outsider_id, project_id, "read"))
|
||||
self.assertFalse(
|
||||
self.store.check_permission(outsider_id, project_id, "create_task"))
|
||||
with self.assertRaises(PermissionError):
|
||||
self.store.require_permission(outsider_id, project_id, "read")
|
||||
|
||||
def test_permission_grants(self):
|
||||
project_id, owner_id, viewer_id, contributor_id, maintainer_id, _ = (
|
||||
self._project_with_roles()
|
||||
)
|
||||
self.assertTrue(self.store.check_permission(viewer_id, project_id, "read"))
|
||||
self.assertTrue(
|
||||
self.store.check_permission(contributor_id, project_id, "create_task"))
|
||||
self.assertTrue(
|
||||
self.store.check_permission(contributor_id, project_id, "read"))
|
||||
self.assertTrue(
|
||||
self.store.check_permission(maintainer_id, project_id, "approve_merge"))
|
||||
self.assertTrue(
|
||||
self.store.check_permission(owner_id, project_id, "manage_project"))
|
||||
self.assertFalse(
|
||||
self.store.check_permission(maintainer_id, project_id, "manage_project"))
|
||||
self.assertEqual(
|
||||
ROLES, ["viewer", "contributor", "maintainer", "owner"])
|
||||
|
||||
def test_owner_can_set_member_maintainer_cannot(self):
|
||||
project_id, owner_id, _v, _c, maintainer_id, _o = (
|
||||
self._project_with_roles()
|
||||
)
|
||||
target_id, _ = self._user("target")
|
||||
# owner 改成员成功
|
||||
self.store.set_member(project_id, owner_id, target_id, "viewer")
|
||||
self.assertTrue(self.store.check_permission(target_id, project_id, "read"))
|
||||
self.assertFalse(
|
||||
self.store.check_permission(target_id, project_id, "create_task"))
|
||||
# maintainer 改成员被拒,且原角色不变
|
||||
with self.assertRaises(PermissionError):
|
||||
self.store.set_member(
|
||||
project_id, maintainer_id, target_id, "contributor")
|
||||
self.assertFalse(
|
||||
self.store.check_permission(target_id, project_id, "create_task"))
|
||||
|
||||
def test_audit_append_only(self):
|
||||
project_id, owner_id, _v, _c, _m, _o = self._project_with_roles()
|
||||
self.assertFalse(hasattr(AuthStore, "delete_audit"))
|
||||
self.assertFalse(hasattr(AuthStore, "update_audit"))
|
||||
before = len(self.store.list_audit(project_id))
|
||||
self.store.audit(owner_id, project_id, "create_task", "task-1")
|
||||
self.store.audit(owner_id, project_id, "approve_merge", "task-1")
|
||||
entries = self.store.list_audit(project_id)
|
||||
self.assertEqual(len(entries), before + 2)
|
||||
# 最新两条在前(倒序),且均为本次追加的动作
|
||||
actions = [e["action"] for e in entries[:2]]
|
||||
self.assertEqual(actions, ["approve_merge", "create_task"])
|
||||
# 审计条目为 dict 且含关键字段
|
||||
self.assertIn("user_id", entries[0])
|
||||
self.assertIn("created_at", entries[0])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,182 @@
|
||||
"""Tests for cloud/memory.py (MemoryStore) and cloud/kb.py (KBStore).
|
||||
|
||||
Run from ~/workspace/codex-app-server:
|
||||
python3 -m unittest discover -s cloud/tests -p "test_memory_kb.py" -v
|
||||
|
||||
No network, no real model turns, no secrets. Each test uses a fresh
|
||||
temporary SQLite file.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.normpath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from kb import KBStore # noqa: E402
|
||||
from memory import MemoryStore # noqa: E402
|
||||
|
||||
|
||||
def _tmp_db():
|
||||
fd, path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
os.unlink(path) # let sqlite create it fresh
|
||||
return path
|
||||
|
||||
|
||||
class MemoryStoreTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.db = _tmp_db()
|
||||
self.store = MemoryStore(self.db)
|
||||
|
||||
def tearDown(self):
|
||||
try:
|
||||
self.store.conn.close()
|
||||
finally:
|
||||
if os.path.exists(self.db):
|
||||
os.unlink(self.db)
|
||||
|
||||
def test_add_then_list(self):
|
||||
sid = self.store.add_summary(
|
||||
"proj-a", "thread-1", "add login page", "login page done",
|
||||
files_changed=["app.py"], open_issues=["tests pending"],
|
||||
)
|
||||
self.assertIsInstance(sid, int)
|
||||
rows = self.store.list_summaries("proj-a")
|
||||
self.assertEqual(len(rows), 1)
|
||||
self.assertEqual(rows[0]["thread_id"], "thread-1")
|
||||
self.assertEqual(rows[0]["goal"], "add login page")
|
||||
self.assertEqual(rows[0]["files_changed"], ["app.py"])
|
||||
self.assertEqual(rows[0]["open_issues"], ["tests pending"])
|
||||
self.assertEqual(rows[0]["status"], "active")
|
||||
|
||||
def test_deprecate_hidden_by_default_visible_with_flag(self):
|
||||
sid = self.store.add_summary("proj-a", "t1", "g1", "c1")
|
||||
self.store.add_summary("proj-a", "t2", "g2", "c2")
|
||||
self.store.deprecate(sid)
|
||||
default_rows = self.store.list_summaries("proj-a")
|
||||
self.assertEqual([r["thread_id"] for r in default_rows], ["t2"])
|
||||
all_rows = self.store.list_summaries("proj-a", include_deprecated=True)
|
||||
self.assertEqual(len(all_rows), 2)
|
||||
deprecated = [r for r in all_rows if r["id"] == sid]
|
||||
self.assertEqual(len(deprecated), 1)
|
||||
self.assertEqual(deprecated[0]["status"], "deprecated")
|
||||
|
||||
def test_delete_removes(self):
|
||||
sid = self.store.add_summary("proj-a", "t1", "g1", "c1")
|
||||
self.store.delete(sid)
|
||||
self.assertEqual(self.store.list_summaries("proj-a"), [])
|
||||
self.assertEqual(
|
||||
self.store.list_summaries("proj-a", include_deprecated=True), []
|
||||
)
|
||||
|
||||
def test_build_prompt_context_top_k_and_excludes_deprecated(self):
|
||||
# 7 active summaries; only the newest 5 may appear.
|
||||
for i in range(7):
|
||||
self.store.add_summary("proj-a", "t%d" % i, "goal-%d" % i, "done-%d" % i)
|
||||
ctx = self.store.build_prompt_context("proj-a", top_k=5)
|
||||
self.assertEqual(ctx.count("\n") + 1, 5)
|
||||
for i in range(2, 7):
|
||||
self.assertIn("goal-%d" % i, ctx)
|
||||
for i in range(2):
|
||||
self.assertNotIn("goal-%d" % i, ctx)
|
||||
# Format spot-check on one line.
|
||||
self.assertIn("- [t6] 目标:goal-6;结论:done-6;", ctx)
|
||||
# A deprecated newest summary must not appear even within top_k.
|
||||
sid = self.store.add_summary("proj-a", "t-new", "goal-new", "done-new")
|
||||
self.store.deprecate(sid)
|
||||
ctx2 = self.store.build_prompt_context("proj-a", top_k=5)
|
||||
self.assertNotIn("goal-new", ctx2)
|
||||
self.assertIn("goal-6", ctx2)
|
||||
|
||||
def test_build_prompt_context_empty(self):
|
||||
self.assertEqual(self.store.build_prompt_context("proj-empty"), "")
|
||||
|
||||
|
||||
class KBStoreTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.db = _tmp_db()
|
||||
self.store = KBStore(self.db)
|
||||
|
||||
def tearDown(self):
|
||||
try:
|
||||
self.store.conn.close()
|
||||
finally:
|
||||
if os.path.exists(self.db):
|
||||
os.unlink(self.db)
|
||||
|
||||
def test_search_finds_file_by_keyword(self):
|
||||
self.store.index_file(
|
||||
"proj-a", "docs/auth.md",
|
||||
"authentication flow uses token refresh and session cookies",
|
||||
kind="doc",
|
||||
)
|
||||
self.store.index_file(
|
||||
"proj-a", "docs/other.md", "unrelated gardening notes", kind="doc"
|
||||
)
|
||||
results = self.store.search("proj-a", "authentication")
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(results[0]["path"], "docs/auth.md")
|
||||
self.assertEqual(results[0]["kind"], "doc")
|
||||
self.assertIn("authentication", results[0]["snippet"])
|
||||
self.assertLessEqual(len(results[0]["snippet"]), 200)
|
||||
|
||||
def test_search_project_isolation(self):
|
||||
self.store.index_file(
|
||||
"proj-a", "docs/secret.md", "zebra secret keyword", kind="doc"
|
||||
)
|
||||
self.store.index_file(
|
||||
"proj-b", "docs/plain.md", "plain notes with no keyword", kind="doc"
|
||||
)
|
||||
# A's doc must never leak into B's search, and vice versa.
|
||||
self.assertEqual(self.store.search("proj-b", "zebra"), [])
|
||||
self.assertEqual(self.store.search("proj-a", "zebra")[0]["path"], "docs/secret.md")
|
||||
|
||||
def test_reindex_same_path_no_duplicates(self):
|
||||
self.store.index_file("proj-a", "README.md", "version one content", kind="readme")
|
||||
self.store.index_file("proj-a", "README.md", "version two content", kind="readme")
|
||||
docs = self.store.list_docs("proj-a")
|
||||
self.assertEqual(len(docs), 1)
|
||||
self.assertEqual(docs[0]["path"], "README.md")
|
||||
self.assertEqual(docs[0]["content"], "version two content")
|
||||
# Search must also see only the new version, once.
|
||||
self.assertEqual(len(self.store.search("proj-a", "version")), 1)
|
||||
self.assertEqual(self.store.search("proj-a", "version")[0]["snippet"],
|
||||
"version two content")
|
||||
self.assertEqual(self.store.search("proj-a", "one"), [])
|
||||
|
||||
def test_index_task_and_list_docs(self):
|
||||
tid = self.store.index_task(
|
||||
"proj-a", "task-42", "fix the login bug", "fixed by resetting tokens"
|
||||
)
|
||||
self.assertIsInstance(tid, int)
|
||||
docs = self.store.list_docs("proj-a")
|
||||
self.assertEqual(len(docs), 1)
|
||||
self.assertEqual(docs[0]["path"], "task:task-42")
|
||||
self.assertEqual(docs[0]["kind"], "task")
|
||||
results = self.store.search("proj-a", "login")
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(results[0]["path"], "task:task-42")
|
||||
|
||||
def test_search_chinese_keyword(self):
|
||||
# FTS5 unicode61 does not tokenize CJK words, so a verbatim
|
||||
# Chinese keyword must still be found via the LIKE fallback.
|
||||
self.store.index_file(
|
||||
"proj-a", "docs/login.md", "用户登录流程使用令牌刷新", kind="doc"
|
||||
)
|
||||
results = self.store.search("proj-a", "登录")
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(results[0]["path"], "docs/login.md")
|
||||
|
||||
def test_search_top_k_and_empty_query(self):
|
||||
for i in range(4):
|
||||
self.store.index_file(
|
||||
"proj-a", "doc%d.md" % i, "shared keyword doc %d" % i, kind="doc"
|
||||
)
|
||||
self.assertEqual(len(self.store.search("proj-a", "shared", top_k=2)), 2)
|
||||
self.assertEqual(self.store.search("proj-a", ""), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Tests for cloud.models (fallback routing) and cloud.cost.
|
||||
|
||||
No real model turns, no network, no secrets: usage clients are fakes
|
||||
and the CostStore uses an in-memory SQLite database.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from datetime import datetime, timezone
|
||||
|
||||
# cloud/ has no __init__.py by design (the coordinator creates it), so
|
||||
# make the project root importable regardless of the discover top level.
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(
|
||||
os.path.dirname(os.path.abspath(__file__)))))
|
||||
|
||||
from cloud.cost import CostStore, estimate_tokens, fetch_account_usage
|
||||
from cloud.models import FALLBACKABLE, ModelRouter, is_fallbackable, record_fallback
|
||||
|
||||
|
||||
class FakeError(Exception):
|
||||
def __init__(self, status=None, code=None):
|
||||
super().__init__(status or code)
|
||||
self.status = status
|
||||
self.code = code
|
||||
|
||||
|
||||
class TestIsFallbackable(unittest.TestCase):
|
||||
def test_fallbackable_status_codes(self):
|
||||
for code in (404, 429, 500, 502, 503, 504):
|
||||
self.assertTrue(is_fallbackable(code), code)
|
||||
self.assertIn(404, FALLBACKABLE)
|
||||
self.assertIn(429, FALLBACKABLE)
|
||||
self.assertIn("timeout", FALLBACKABLE)
|
||||
|
||||
def test_timeout_string(self):
|
||||
self.assertTrue(is_fallbackable("timeout"))
|
||||
|
||||
def test_not_fallbackable(self):
|
||||
for code in (400, 401, 403, 200):
|
||||
self.assertFalse(is_fallbackable(code), code)
|
||||
self.assertFalse(is_fallbackable("other"))
|
||||
|
||||
def test_object_with_status_or_code(self):
|
||||
self.assertTrue(is_fallbackable(FakeError(status=429)))
|
||||
self.assertTrue(is_fallbackable(FakeError(code=503)))
|
||||
self.assertTrue(is_fallbackable(FakeError(status="timeout")))
|
||||
self.assertFalse(is_fallbackable(FakeError(status=401)))
|
||||
self.assertFalse(is_fallbackable(FakeError(code=400)))
|
||||
|
||||
|
||||
class TestModelRouter(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.router = ModelRouter(["gpt-6.1-sol", "gpt-6-sol"])
|
||||
|
||||
def test_empty_priority_list_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
ModelRouter([])
|
||||
|
||||
def test_first(self):
|
||||
self.assertEqual(self.router.first(), "gpt-6.1-sol")
|
||||
|
||||
def test_next_after_in_priority_order(self):
|
||||
self.assertEqual(
|
||||
self.router.next_after("gpt-6.1-sol", 404), "gpt-6-sol")
|
||||
self.assertEqual(
|
||||
self.router.next_after("gpt-6.1-sol", "timeout"), "gpt-6-sol")
|
||||
|
||||
def test_next_after_exhausted_returns_none(self):
|
||||
self.assertIsNone(self.router.next_after("gpt-6-sol", 500))
|
||||
self.assertIsNone(self.router.next_after("gpt-6-sol", 429))
|
||||
|
||||
def test_next_after_non_fallbackable_returns_none(self):
|
||||
self.assertIsNone(self.router.next_after("gpt-6.1-sol", 401))
|
||||
self.assertIsNone(self.router.next_after("gpt-6.1-sol", 400))
|
||||
|
||||
def test_route_walks_attempts(self):
|
||||
self.assertEqual(
|
||||
self.router.route([("gpt-6.1-sol", 404)]), "gpt-6-sol")
|
||||
self.assertIsNone(self.router.route(
|
||||
[("gpt-6.1-sol", 404), ("gpt-6-sol", 500)]))
|
||||
self.assertIsNone(self.router.route([("gpt-6.1-sol", 401)]))
|
||||
self.assertEqual(self.router.route([]), "gpt-6.1-sol")
|
||||
|
||||
|
||||
class TestRecordFallback(unittest.TestCase):
|
||||
def test_event_recorded_on_task_record(self):
|
||||
task = {"id": "task-1"}
|
||||
event = record_fallback(task, "gpt-6.1-sol", "gpt-6-sol", "429")
|
||||
self.assertEqual(event, {
|
||||
"type": "model_fallback",
|
||||
"task_id": "task-1",
|
||||
"from_model": "gpt-6.1-sol",
|
||||
"to_model": "gpt-6-sol",
|
||||
"reason": "429",
|
||||
})
|
||||
self.assertIn(event, task["events"])
|
||||
|
||||
def test_event_appended_to_existing_events(self):
|
||||
task = {"id": "task-2", "events": [{"type": "started"}]}
|
||||
record_fallback(task, "a", "b", "500")
|
||||
self.assertEqual(len(task["events"]), 2)
|
||||
self.assertEqual(task["events"][1]["type"], "model_fallback")
|
||||
|
||||
|
||||
class TestFetchAccountUsage(unittest.TestCase):
|
||||
def test_success_returns_available_true(self):
|
||||
data = {"balance": 123, "currency": "USD"}
|
||||
result = fetch_account_usage(lambda: data)
|
||||
self.assertEqual(result, {"available": True, "data": data})
|
||||
|
||||
def test_raising_client_returns_unavailable_not_raise(self):
|
||||
def boom():
|
||||
raise RuntimeError("no route")
|
||||
self.assertEqual(
|
||||
fetch_account_usage(boom), {"available": False, "data": None})
|
||||
|
||||
def test_none_and_empty_dict_return_unavailable(self):
|
||||
self.assertEqual(
|
||||
fetch_account_usage(lambda: None),
|
||||
{"available": False, "data": None})
|
||||
self.assertEqual(
|
||||
fetch_account_usage(lambda: {}),
|
||||
{"available": False, "data": None})
|
||||
|
||||
|
||||
class TestEstimateTokens(unittest.TestCase):
|
||||
def test_chars_divided_by_four(self):
|
||||
self.assertEqual(estimate_tokens("abcd", "efgh"), 2)
|
||||
self.assertEqual(estimate_tokens("a" * 10, "b" * 7), 4) # 17 // 4
|
||||
self.assertEqual(estimate_tokens("", ""), 0)
|
||||
|
||||
|
||||
class TestCostStore(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.store = CostStore(":memory:")
|
||||
|
||||
def test_record_computes_total_and_day(self):
|
||||
row_id = self.store.record(
|
||||
"t1", "p1", "gpt-6.1-sol", input_tokens=100, output_tokens=50,
|
||||
usd_micros=2000, source="actual", fallback_count=1)
|
||||
self.assertIsInstance(row_id, int)
|
||||
report = self.store.report_by_task("t1")
|
||||
self.assertEqual(report["total_tokens"], 150)
|
||||
self.assertEqual(report["usd_micros"], 2000)
|
||||
self.assertEqual(report["turns"], 1)
|
||||
self.assertEqual(report["fallback_count"], 1)
|
||||
self.assertEqual(report["actual_tokens"], 150)
|
||||
self.assertEqual(report["estimated_tokens"], 0)
|
||||
|
||||
def test_invalid_source_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.record("t1", "p1", "m", source="guess")
|
||||
|
||||
def test_actual_and_estimated_kept_separate(self):
|
||||
self.store.record("t1", "p1", "m", input_tokens=100,
|
||||
output_tokens=0, source="actual")
|
||||
self.store.record_estimated("t1", "p1", "m",
|
||||
"x" * 40, "y" * 40) # 80 // 4 = 20
|
||||
report = self.store.report_by_task("t1")
|
||||
self.assertEqual(report["total_tokens"], 120)
|
||||
self.assertEqual(report["actual_tokens"], 100)
|
||||
self.assertEqual(report["estimated_tokens"], 20)
|
||||
self.assertEqual(report["turns"], 2)
|
||||
|
||||
def test_record_estimated_shape(self):
|
||||
self.store.record_estimated("t9", "p9", "m", "abcd", "efgh")
|
||||
report = self.store.report_by_task("t9")
|
||||
self.assertEqual(report["estimated_tokens"], 2)
|
||||
self.assertEqual(report["actual_tokens"], 0)
|
||||
self.assertEqual(report["usd_micros"], 0)
|
||||
|
||||
def test_report_by_project_and_day(self):
|
||||
self.store.record("t1", "p1", "m", input_tokens=10, output_tokens=0)
|
||||
self.store.record("t2", "p1", "m", input_tokens=30, output_tokens=0)
|
||||
self.store.record("t3", "p2", "m", input_tokens=999, output_tokens=0)
|
||||
project = self.store.report_by_project("p1")
|
||||
self.assertEqual(project["total_tokens"], 40)
|
||||
self.assertEqual(project["turns"], 2)
|
||||
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||||
day = self.store.report_by_day(today)
|
||||
self.assertEqual(day["total_tokens"], 10 + 30 + 999)
|
||||
self.assertEqual(day["turns"], 3)
|
||||
self.assertEqual(self.store.report_by_day("2000-01-01")["turns"], 0)
|
||||
|
||||
def test_empty_report_is_zeroed(self):
|
||||
report = self.store.report_by_task("missing")
|
||||
self.assertEqual(report, {
|
||||
"total_tokens": 0,
|
||||
"usd_micros": 0,
|
||||
"turns": 0,
|
||||
"fallback_count": 0,
|
||||
"actual_tokens": 0,
|
||||
"estimated_tokens": 0,
|
||||
})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,206 @@
|
||||
"""派发队列与沙箱执行器的单测。
|
||||
|
||||
运行(在 ~/workspace/codex-app-server 下):
|
||||
python3 -m unittest discover -s cloud/tests -p "test_queue_sandbox.py" -v
|
||||
|
||||
说明:cloud/ 没有 __init__.py(协调员统一创建),因此这里把 cloud/
|
||||
目录插到 sys.path 前面再导入 queue / sandbox,与 discover 的
|
||||
top-level 目录无关。
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(
|
||||
os.path.abspath(__file__))))
|
||||
|
||||
from queue import TaskQueue, budget_exceeded, order_tasks, ready_tasks
|
||||
from sandbox import CloudExecutor, DockerExecutor, Executor, LocalExecutor
|
||||
|
||||
|
||||
def _task(task_id, depends_on=None, priority=0, status="queued"):
|
||||
task = {"id": task_id, "priority": priority, "status": status}
|
||||
if depends_on is not None:
|
||||
task["depends_on"] = depends_on
|
||||
return task
|
||||
|
||||
|
||||
class TestOrderTasks(unittest.TestCase):
|
||||
def test_dependency_order(self):
|
||||
# A -> B 时 B 在 A 后
|
||||
ordered = order_tasks([_task("B", depends_on=["A"]),
|
||||
_task("A")])
|
||||
self.assertEqual([t["id"] for t in ordered], ["A", "B"])
|
||||
|
||||
def test_dependency_chain(self):
|
||||
ordered = order_tasks([_task("C", depends_on=["B"]),
|
||||
_task("B", depends_on=["A"]),
|
||||
_task("A")])
|
||||
self.assertEqual([t["id"] for t in ordered], ["A", "B", "C"])
|
||||
|
||||
def test_priority_within_wave(self):
|
||||
# 同一波(均无依赖)按 priority 大的先
|
||||
ordered = order_tasks([_task("low", priority=1),
|
||||
_task("high", priority=10),
|
||||
_task("mid", priority=5)])
|
||||
self.assertEqual([t["id"] for t in ordered],
|
||||
["high", "mid", "low"])
|
||||
|
||||
def test_priority_does_not_jump_dependency(self):
|
||||
# 高优先级也不能越过依赖:B 依赖 A,即使 B 优先级更高也排 A 后
|
||||
ordered = order_tasks([_task("B", depends_on=["A"], priority=99),
|
||||
_task("A", priority=0)])
|
||||
self.assertEqual([t["id"] for t in ordered], ["A", "B"])
|
||||
|
||||
def test_missing_dependency_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
order_tasks([_task("B", depends_on=["NOPE"])])
|
||||
|
||||
def test_cycle_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
order_tasks([_task("A", depends_on=["B"]),
|
||||
_task("B", depends_on=["A"])])
|
||||
|
||||
def test_self_dependency_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
order_tasks([_task("A", depends_on=["A"])])
|
||||
|
||||
def test_defaults(self):
|
||||
# depends_on / priority 缺省时按 0 处理、不报错
|
||||
ordered = order_tasks([{"id": "A"}, {"id": "B"}])
|
||||
self.assertEqual([t["id"] for t in ordered], ["A", "B"])
|
||||
|
||||
|
||||
class TestReadyTasks(unittest.TestCase):
|
||||
def test_only_dependency_satisfied_queued(self):
|
||||
tasks = [_task("A"),
|
||||
_task("B", depends_on=["A"]),
|
||||
_task("C", depends_on=["A"], status="running"),
|
||||
_task("D", depends_on=["MISSING"])]
|
||||
ready = ready_tasks(tasks, completed_ids=[])
|
||||
self.assertEqual([t["id"] for t in ready], ["A"])
|
||||
|
||||
# A 已完成(状态随之改为 completed)后,B 依赖满足且 queued;
|
||||
# C 是 running 不返回;D 依赖未满足不返回
|
||||
done_tasks = [_task("A", status="completed"),
|
||||
_task("B", depends_on=["A"]),
|
||||
_task("C", depends_on=["A"], status="running"),
|
||||
_task("D", depends_on=["MISSING"])]
|
||||
ready = ready_tasks(done_tasks, completed_ids=["A"])
|
||||
self.assertEqual([t["id"] for t in ready], ["B"])
|
||||
|
||||
def test_ready_ordered_by_priority(self):
|
||||
tasks = [_task("low", priority=1), _task("high", priority=9)]
|
||||
ready = ready_tasks(tasks, completed_ids=[])
|
||||
self.assertEqual([t["id"] for t in ready], ["high", "low"])
|
||||
|
||||
|
||||
class TestBudgetExceeded(unittest.TestCase):
|
||||
def test_below_limits(self):
|
||||
self.assertFalse(budget_exceeded(5, 10, 0.5, 1.0))
|
||||
|
||||
def test_at_limit_counts_as_exceeded(self):
|
||||
self.assertTrue(budget_exceeded(10, 10, 0.5, 1.0))
|
||||
self.assertTrue(budget_exceeded(5, 10, 1.0, 1.0))
|
||||
|
||||
def test_none_limit_means_unlimited(self):
|
||||
self.assertFalse(budget_exceeded(10 ** 9, None, 10 ** 9, None))
|
||||
self.assertTrue(budget_exceeded(10 ** 9, None, 10 ** 9, 1.0))
|
||||
|
||||
|
||||
class TestTaskQueue(unittest.TestCase):
|
||||
def test_next_ready_flow(self):
|
||||
queue = TaskQueue()
|
||||
queue.add(_task("A"))
|
||||
queue.add(_task("B", depends_on=["A"]))
|
||||
self.assertEqual(queue.next_ready([])["id"], "A")
|
||||
queue.mark_running("A")
|
||||
# A 在 running,不再是可派任务;B 依赖未完成,也不可派
|
||||
self.assertIsNone(queue.next_ready([]))
|
||||
queue.mark_done("A", "completed")
|
||||
self.assertEqual(queue.next_ready(["A"])["id"], "B")
|
||||
|
||||
def test_mark_done_validates_status(self):
|
||||
queue = TaskQueue([_task("A")])
|
||||
with self.assertRaises(ValueError):
|
||||
queue.mark_done("A", "bogus-status")
|
||||
with self.assertRaises(ValueError):
|
||||
queue.add(_task("A")) # 重复 id
|
||||
|
||||
def test_all_host_statuses_accepted(self):
|
||||
queue = TaskQueue()
|
||||
for status in ("running", "completed", "failed",
|
||||
"interrupted", "merged", "merge-failed"):
|
||||
queue.add(_task("t-" + status, status=status))
|
||||
# 非 queued 任务一律不可派
|
||||
self.assertIsNone(queue.next_ready([]))
|
||||
|
||||
|
||||
class TestExecutors(unittest.TestCase):
|
||||
def test_base_executor_must_implement(self):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
Executor().execute({"id": "x"})
|
||||
|
||||
def test_local_executor_success(self):
|
||||
seen = []
|
||||
|
||||
def fake_runner(cmd):
|
||||
seen.append(cmd)
|
||||
return SimpleNamespace(returncode=0, stdout="ok-status")
|
||||
|
||||
executor = LocalExecutor("/fake/host.py", config_path="/fake/cfg.json",
|
||||
runner=fake_runner)
|
||||
result = executor.execute({"id": "t1", "project_id": "p",
|
||||
"repo": "/r", "prompt": "hi"})
|
||||
self.assertEqual(result["executor"], "local")
|
||||
self.assertEqual(result["task_id"], "t1")
|
||||
self.assertEqual(result["status"], "completed")
|
||||
self.assertEqual(result["stdout"], "ok-status")
|
||||
# 只发只读 status 命令,不发 dispatch/run(不直接跑模型)
|
||||
self.assertIn("status", seen[0])
|
||||
self.assertNotIn("dispatch", seen[0])
|
||||
self.assertNotIn("run", seen[0])
|
||||
self.assertIn("--config", seen[0])
|
||||
|
||||
def test_local_executor_failure(self):
|
||||
def fake_runner(cmd):
|
||||
return SimpleNamespace(returncode=1, stdout="boom")
|
||||
|
||||
result = LocalExecutor("/fake/host.py",
|
||||
runner=fake_runner).execute({"id": "t2"})
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertEqual(result["stdout"], "boom")
|
||||
|
||||
def test_local_executor_truncates_stdout(self):
|
||||
def fake_runner(cmd):
|
||||
return SimpleNamespace(returncode=0, stdout="x" * 5000)
|
||||
|
||||
result = LocalExecutor("/fake/host.py",
|
||||
runner=fake_runner).execute({"id": "t3"})
|
||||
self.assertEqual(len(result["stdout"]), 2000)
|
||||
|
||||
def test_docker_executor_no_docker_raises_not_local(self):
|
||||
with mock.patch("sandbox.shutil.which", return_value=None):
|
||||
with self.assertRaises(RuntimeError) as ctx:
|
||||
DockerExecutor().execute({"id": "t1"})
|
||||
self.assertIn("docker not available", str(ctx.exception))
|
||||
# 报错而不是返回 local 结果:绝不静默降级为本地执行
|
||||
|
||||
def test_docker_executor_with_docker_is_placeholder_failure(self):
|
||||
with mock.patch("sandbox.shutil.which", return_value="/usr/bin/docker"):
|
||||
result = DockerExecutor().execute({"id": "t1"})
|
||||
self.assertEqual(result["executor"], "docker")
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertNotEqual(result["executor"], "local")
|
||||
|
||||
def test_cloud_executor_raises_with_cloud_account(self):
|
||||
with self.assertRaises(RuntimeError) as ctx:
|
||||
CloudExecutor().execute({"id": "t1"})
|
||||
self.assertIn("cloud account", str(ctx.exception))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,311 @@
|
||||
"""控制面 server 集成测试:对本地临时端口做真实 HTTP 请求。
|
||||
|
||||
只用标准库(http.client),不跑真实模型回合,不联网连真实
|
||||
GitHub/Slack(webhook 测试用的 secret 均为单测内临时值)。
|
||||
每个测试用独立的临时 SQLite 文件与临时端口,跑完关闭 server。
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
try:
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.server import create_server, sweep_stale_tasks
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.server import create_server, sweep_stale_tasks
|
||||
|
||||
|
||||
class ServerTestCase(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.viewer_id, self.viewer_key = self.auth.create_user("viewer")
|
||||
self.outsider_id, self.outsider_key = self.auth.create_user("outsider")
|
||||
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()
|
||||
# owner 建一个项目,viewer 设为该项目的 viewer 角色。
|
||||
status, self.project = self._req(
|
||||
"POST", "/projects", {"name": "demo"}, self.owner_key)
|
||||
self.assertEqual(status, 201)
|
||||
self.auth.set_member(
|
||||
self.project["id"], self.owner_id, self.viewer_id, "viewer")
|
||||
|
||||
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()
|
||||
|
||||
# -- HTTP 工具 -----------------------------------------------------
|
||||
def _req(self, method, path, body=None, key=None, headers=None,
|
||||
raw_body=None):
|
||||
conn = http.client.HTTPConnection(
|
||||
"127.0.0.1", self.port, timeout=10)
|
||||
hdrs = dict(headers or {})
|
||||
if raw_body is not None:
|
||||
payload = raw_body
|
||||
elif body is not None:
|
||||
payload = json.dumps(body).encode("utf-8")
|
||||
hdrs.setdefault("Content-Type", "application/json")
|
||||
else:
|
||||
payload = None
|
||||
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 _create_task(self, key, **extra):
|
||||
body = {"project_id": self.project["id"], "prompt": "做一个加法函数"}
|
||||
body.update(extra)
|
||||
status, task = self._req("POST", "/tasks", body, key)
|
||||
self.assertEqual(status, 201, task)
|
||||
return task
|
||||
|
||||
# -- 健康检查与认证 -------------------------------------------------
|
||||
def test_health_no_auth(self):
|
||||
status, body = self._req("GET", "/health")
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(
|
||||
body, {"ok": True, "service": "cloud-control-plane"})
|
||||
|
||||
def test_no_key_401(self):
|
||||
status, body = self._req("GET", "/projects")
|
||||
self.assertEqual(status, 401)
|
||||
self.assertEqual(body, {"error": "unauthorized"})
|
||||
|
||||
def test_wrong_key_401(self):
|
||||
status, body = self._req("GET", "/projects", key="wrong-key")
|
||||
self.assertEqual(status, 401)
|
||||
self.assertEqual(body, {"error": "unauthorized"})
|
||||
|
||||
# -- 项目可见性与权限 -----------------------------------------------
|
||||
def test_projects_only_member_projects(self):
|
||||
# outsider 建自己的项目;双方互不可见对方的项目。
|
||||
status, other = self._req(
|
||||
"POST", "/projects", {"name": "other"}, self.outsider_key)
|
||||
self.assertEqual(status, 201)
|
||||
status, body = self._req("GET", "/projects", key=self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
ids = [p["id"] for p in body["projects"]]
|
||||
self.assertIn(self.project["id"], ids)
|
||||
self.assertNotIn(other["id"], ids)
|
||||
status, body = self._req("GET", "/projects", key=self.outsider_key)
|
||||
self.assertEqual(status, 200)
|
||||
ids = [p["id"] for p in body["projects"]]
|
||||
self.assertIn(other["id"], ids)
|
||||
self.assertNotIn(self.project["id"], ids)
|
||||
|
||||
def test_nonmember_read_project_tasks_403(self):
|
||||
status, body = self._req(
|
||||
"GET", "/tasks?project_id=%s" % self.project["id"],
|
||||
key=self.outsider_key)
|
||||
self.assertEqual(status, 403)
|
||||
|
||||
def test_viewer_create_task_403(self):
|
||||
status, body = self._req(
|
||||
"POST", "/tasks",
|
||||
{"project_id": self.project["id"], "prompt": "x"},
|
||||
self.viewer_key)
|
||||
self.assertEqual(status, 403)
|
||||
|
||||
def test_viewer_can_read_tasks(self):
|
||||
self._create_task(self.owner_key)
|
||||
status, body = self._req(
|
||||
"GET", "/tasks?project_id=%s" % self.project["id"],
|
||||
key=self.viewer_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(len(body["tasks"]), 1)
|
||||
|
||||
# -- 任务 CRUD、幂等、取消、心跳、回收 -------------------------------
|
||||
def test_create_and_get_task(self):
|
||||
task = self._create_task(self.owner_key, priority=3)
|
||||
self.assertEqual(task["status"], "queued")
|
||||
self.assertEqual(task["priority"], 3)
|
||||
status, got = self._req(
|
||||
"GET", "/tasks/%s" % task["id"], key=self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(got["id"], task["id"])
|
||||
status, _ = self._req("GET", "/tasks/nope", key=self.owner_key)
|
||||
self.assertEqual(status, 404)
|
||||
status, _ = self._req(
|
||||
"GET", "/tasks/%s" % task["id"], key=self.outsider_key)
|
||||
self.assertEqual(status, 403)
|
||||
|
||||
def test_idempotent_key_returns_same_task(self):
|
||||
first = self._create_task(
|
||||
self.owner_key, idempotency_key="key-1")
|
||||
status, second = self._req(
|
||||
"POST", "/tasks",
|
||||
{"project_id": self.project["id"], "prompt": "另一个 prompt",
|
||||
"idempotency_key": "key-1"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(second["id"], first["id"])
|
||||
self.assertEqual(second["prompt"], first["prompt"])
|
||||
status, body = self._req(
|
||||
"GET", "/tasks?project_id=%s" % self.project["id"],
|
||||
key=self.owner_key)
|
||||
self.assertEqual(len(body["tasks"]), 1)
|
||||
|
||||
def test_cancel_running_then_terminal_cancel_400(self):
|
||||
task = self._create_task(self.owner_key)
|
||||
status, running = self._req(
|
||||
"POST", "/tasks/%s/heartbeat" % task["id"], {}, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(running["status"], "running")
|
||||
status, cancelled = self._req(
|
||||
"POST", "/tasks/%s/cancel" % task["id"], {}, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(cancelled["status"], "interrupted")
|
||||
status, body = self._req(
|
||||
"POST", "/tasks/%s/cancel" % task["id"], {}, self.owner_key)
|
||||
self.assertEqual(status, 400)
|
||||
|
||||
def test_cancel_queued_task(self):
|
||||
task = self._create_task(self.owner_key)
|
||||
status, body = self._req(
|
||||
"POST", "/tasks/%s/cancel" % task["id"], {}, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(body["status"], "interrupted")
|
||||
|
||||
def test_heartbeat_after_sweep_keeps_fresh_task(self):
|
||||
task = self._create_task(self.owner_key)
|
||||
self._req("POST", "/tasks/%s/heartbeat" % task["id"], {},
|
||||
self.owner_key)
|
||||
status, body = self._req(
|
||||
"POST", "/admin/sweep", {}, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertNotIn(task["id"], body["swept"])
|
||||
status, got = self._req(
|
||||
"GET", "/tasks/%s" % task["id"], key=self.owner_key)
|
||||
self.assertEqual(got["status"], "running")
|
||||
|
||||
def test_sweep_stale_with_threshold_zero(self):
|
||||
task = self._create_task(self.owner_key)
|
||||
self._req("POST", "/tasks/%s/heartbeat" % task["id"], {},
|
||||
self.owner_key)
|
||||
swept = sweep_stale_tasks(self.db_path, threshold_seconds=0)
|
||||
self.assertIn(task["id"], swept)
|
||||
status, got = self._req(
|
||||
"GET", "/tasks/%s" % task["id"], key=self.owner_key)
|
||||
self.assertEqual(got["status"], "interrupted")
|
||||
|
||||
def test_sweep_stale_manually_aged_heartbeat(self):
|
||||
task = self._create_task(self.owner_key)
|
||||
self._req("POST", "/tasks/%s/heartbeat" % task["id"], {},
|
||||
self.owner_key)
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"UPDATE tasks SET last_heartbeat = ? WHERE id = ?",
|
||||
(time.time() - 3600, task["id"]))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
status, body = self._req(
|
||||
"POST", "/admin/sweep", {}, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertIn(task["id"], body["swept"])
|
||||
status, got = self._req(
|
||||
"GET", "/tasks/%s" % task["id"], key=self.owner_key)
|
||||
self.assertEqual(got["status"], "interrupted")
|
||||
|
||||
def test_sweep_skips_queued_task(self):
|
||||
task = self._create_task(self.owner_key) # 仍是 queued
|
||||
swept = sweep_stale_tasks(self.db_path, threshold_seconds=0)
|
||||
self.assertNotIn(task["id"], swept)
|
||||
|
||||
# -- webhooks(验签,解析但不执行)----------------------------------
|
||||
def test_github_webhook_secret_missing_503(self):
|
||||
with mock.patch.dict(os.environ, {}, clear=False):
|
||||
os.environ.pop("GITHUB_WEBHOOK_SECRET", None)
|
||||
status, _ = self._req(
|
||||
"POST", "/webhooks/github", {"action": "created"},
|
||||
key=None)
|
||||
self.assertEqual(status, 503)
|
||||
|
||||
def test_github_webhook_bad_signature_401(self):
|
||||
payload = {"action": "created", "comment": {"body": "/fix-review"}}
|
||||
raw = json.dumps(payload).encode("utf-8")
|
||||
with mock.patch.dict(
|
||||
os.environ, {"GITHUB_WEBHOOK_SECRET": "test-secret"}):
|
||||
status, body = self._req(
|
||||
"POST", "/webhooks/github", raw_body=raw,
|
||||
headers={"Content-Type": "application/json",
|
||||
"X-Hub-Signature-256": "sha256=deadbeef"})
|
||||
self.assertEqual(status, 401)
|
||||
|
||||
def test_github_webhook_good_signature_parses(self):
|
||||
payload = {
|
||||
"action": "created",
|
||||
"comment": {"body": "请 /fix-review 处理", "user": {"login": "alice"}},
|
||||
"repository": {"full_name": "org/repo"},
|
||||
"issue": {"number": 7},
|
||||
}
|
||||
raw = json.dumps(payload).encode("utf-8")
|
||||
sig = "sha256=" + hmac.new(
|
||||
b"test-secret", raw, hashlib.sha256).hexdigest()
|
||||
with mock.patch.dict(
|
||||
os.environ, {"GITHUB_WEBHOOK_SECRET": "test-secret"}):
|
||||
status, body = self._req(
|
||||
"POST", "/webhooks/github", raw_body=raw,
|
||||
headers={"Content-Type": "application/json",
|
||||
"X-Hub-Signature-256": sig})
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(body["type"], "fix_review")
|
||||
self.assertEqual(body["pr_number"], 7)
|
||||
|
||||
def test_slack_webhook_bad_and_good_signature(self):
|
||||
raw = b"text=/codex status my-task"
|
||||
ts = str(int(time.time()))
|
||||
good_sig = "v0=" + hmac.new(
|
||||
b"slack-secret",
|
||||
b"v0:" + ts.encode() + b":" + raw, hashlib.sha256).hexdigest()
|
||||
headers = {"Content-Type": "application/x-www-form-urlencoded",
|
||||
"X-Slack-Signature": "v0=deadbeef",
|
||||
"X-Slack-Request-Timestamp": ts}
|
||||
with mock.patch.dict(
|
||||
os.environ, {"SLACK_SIGNING_SECRET": "slack-secret"}):
|
||||
status, _ = self._req(
|
||||
"POST", "/webhooks/slack", raw_body=raw, headers=headers)
|
||||
self.assertEqual(status, 401)
|
||||
headers["X-Slack-Signature"] = good_sig
|
||||
status, body = self._req(
|
||||
"POST", "/webhooks/slack", raw_body=raw, headers=headers)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(body, {"command": "status", "task": "my-task"})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,236 @@
|
||||
"""webhooks.py 单测:签名校验、事件解析、命令派发、用户映射、通知门禁。
|
||||
|
||||
只用标准库,不联网、不连真实 GitHub/Slack、不跑真实模型回合。
|
||||
签名测试用的 secret 均为单测内临时值,不是生产密钥。
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
|
||||
# 本文件位于 cloud/tests/,被测模块位于 cloud/;cloud/ 目前不是包
|
||||
# (协调员统一创建 __init__.py),故把 cloud/ 加到 sys.path 后按
|
||||
# 顶层模块导入 webhooks。
|
||||
sys.path.insert(
|
||||
0, os.path.normpath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
)
|
||||
|
||||
import webhooks # noqa: E402
|
||||
from webhooks import ( # noqa: E402
|
||||
ALLOWED_NOTIFICATION_EVENTS,
|
||||
NotificationService,
|
||||
SlackUserMap,
|
||||
parse_github_pr_comment,
|
||||
parse_slack_command,
|
||||
verify_github_signature,
|
||||
verify_slack_signature,
|
||||
)
|
||||
|
||||
|
||||
def _github_signature(secret, payload):
|
||||
return "sha256=" + hmac.new(
|
||||
secret.encode("utf-8"), payload, hashlib.sha256
|
||||
).hexdigest()
|
||||
|
||||
|
||||
def _slack_signature(secret, timestamp, body):
|
||||
basestring = f"v0:{timestamp}:{body}"
|
||||
return "v0=" + hmac.new(
|
||||
secret.encode("utf-8"), basestring.encode("utf-8"), hashlib.sha256
|
||||
).hexdigest()
|
||||
|
||||
|
||||
class TestGithubSignature(unittest.TestCase):
|
||||
def test_correct_signature_passes(self):
|
||||
secret = "test-github-secret"
|
||||
payload = b'{"action":"created"}'
|
||||
header = _github_signature(secret, payload)
|
||||
self.assertTrue(verify_github_signature(secret, payload, header))
|
||||
|
||||
def test_wrong_signature_rejected(self):
|
||||
secret = "test-github-secret"
|
||||
payload = b'{"action":"created"}'
|
||||
header = _github_signature("wrong-secret", payload)
|
||||
self.assertFalse(verify_github_signature(secret, payload, header))
|
||||
|
||||
def test_tampered_payload_rejected(self):
|
||||
secret = "test-github-secret"
|
||||
header = _github_signature(secret, b'{"action":"created"}')
|
||||
self.assertFalse(
|
||||
verify_github_signature(secret, b'{"action":"deleted"}', header)
|
||||
)
|
||||
|
||||
def test_missing_header_rejected(self):
|
||||
self.assertFalse(
|
||||
verify_github_signature("test-github-secret", b"{}", None)
|
||||
)
|
||||
self.assertFalse(verify_github_signature("test-github-secret", b"{}", ""))
|
||||
|
||||
def test_missing_secret_rejected(self):
|
||||
payload = b"{}"
|
||||
self.assertFalse(verify_github_signature("", payload, "sha256=abc"))
|
||||
self.assertFalse(verify_github_signature(None, payload, "sha256=abc"))
|
||||
|
||||
|
||||
class TestParseGithubPrComment(unittest.TestCase):
|
||||
def _payload(self, body, action="created"):
|
||||
return {
|
||||
"action": action,
|
||||
"repository": {"full_name": "owner/repo"},
|
||||
"issue": {"number": 42},
|
||||
"comment": {"body": body, "user": {"login": "alice"}},
|
||||
}
|
||||
|
||||
def test_fix_review_comment_parsed(self):
|
||||
result = parse_github_pr_comment(
|
||||
self._payload("请帮忙 /fix-review 看一下")
|
||||
)
|
||||
self.assertEqual(
|
||||
result,
|
||||
{
|
||||
"type": "fix_review",
|
||||
"repo": "owner/repo",
|
||||
"pr_number": 42,
|
||||
"comment": "请帮忙 /fix-review 看一下",
|
||||
"author": "alice",
|
||||
},
|
||||
)
|
||||
|
||||
def test_comment_without_trigger_returns_none(self):
|
||||
self.assertIsNone(parse_github_pr_comment(self._payload("好看,LGTM")))
|
||||
|
||||
def test_non_created_action_returns_none(self):
|
||||
self.assertIsNone(
|
||||
parse_github_pr_comment(
|
||||
self._payload("有 /fix-review 但是已编辑", action="edited")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class TestSlackSignature(unittest.TestCase):
|
||||
def test_correct_signature_passes(self):
|
||||
secret = "test-slack-secret"
|
||||
timestamp = str(int(time.time()))
|
||||
body = "command=/codex&text=status"
|
||||
header = _slack_signature(secret, timestamp, body)
|
||||
self.assertTrue(
|
||||
verify_slack_signature(secret, timestamp, body, header)
|
||||
)
|
||||
|
||||
def test_wrong_signature_rejected(self):
|
||||
secret = "test-slack-secret"
|
||||
timestamp = str(int(time.time()))
|
||||
body = "command=/codex&text=status"
|
||||
header = _slack_signature("wrong-secret", timestamp, body)
|
||||
self.assertFalse(
|
||||
verify_slack_signature(secret, timestamp, body, header)
|
||||
)
|
||||
|
||||
def test_expired_timestamp_rejected(self):
|
||||
# 超过 300 秒:用正确密钥签一个 301 秒前的时间戳,仍应拒绝。
|
||||
secret = "test-slack-secret"
|
||||
timestamp = str(int(time.time()) - 301)
|
||||
body = "command=/codex&text=status"
|
||||
header = _slack_signature(secret, timestamp, body)
|
||||
self.assertFalse(
|
||||
verify_slack_signature(secret, timestamp, body, header)
|
||||
)
|
||||
|
||||
def test_missing_params_rejected(self):
|
||||
timestamp = str(int(time.time()))
|
||||
self.assertFalse(
|
||||
verify_slack_signature("", timestamp, "body", "v0=abc")
|
||||
)
|
||||
self.assertFalse(
|
||||
verify_slack_signature("test-slack-secret", timestamp, "body", None)
|
||||
)
|
||||
|
||||
|
||||
class TestParseSlackCommand(unittest.TestCase):
|
||||
def test_dispatch_prompt(self):
|
||||
self.assertEqual(
|
||||
parse_slack_command("/codex 做一个加法函数"),
|
||||
{"command": "codex", "prompt": "做一个加法函数"},
|
||||
)
|
||||
|
||||
def test_dispatch_status(self):
|
||||
self.assertEqual(
|
||||
parse_slack_command("/codex status my-task"),
|
||||
{"command": "status", "task": "my-task"},
|
||||
)
|
||||
|
||||
def test_empty_or_non_codex_returns_none_command(self):
|
||||
self.assertEqual(
|
||||
parse_slack_command(""), {"command": None, "prompt": ""}
|
||||
)
|
||||
self.assertEqual(
|
||||
parse_slack_command("hello"),
|
||||
{"command": None, "prompt": "hello"},
|
||||
)
|
||||
self.assertEqual(
|
||||
parse_slack_command("/other do something"),
|
||||
{"command": None, "prompt": "/other do something"},
|
||||
)
|
||||
|
||||
|
||||
class TestSlackUserMap(unittest.TestCase):
|
||||
def test_map_and_lookup(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
db_path = os.path.join(tmpdir, "slack_users.db")
|
||||
user_map = SlackUserMap(db_path)
|
||||
try:
|
||||
self.assertIsNone(user_map.lookup("U123"))
|
||||
user_map.map("U123", "user-1")
|
||||
self.assertEqual(user_map.lookup("U123"), "user-1")
|
||||
# 覆盖映射
|
||||
user_map.map("U123", "user-2")
|
||||
self.assertEqual(user_map.lookup("U123"), "user-2")
|
||||
# 未映射的 id 仍返回 None
|
||||
self.assertIsNone(user_map.lookup("U999"))
|
||||
finally:
|
||||
user_map.close()
|
||||
|
||||
|
||||
class TestNotificationService(unittest.TestCase):
|
||||
def test_allowed_event_sent_and_senders_called(self):
|
||||
calls = {"webhook": [], "slack": []}
|
||||
service = NotificationService(
|
||||
webhook_sender=lambda payload: calls["webhook"].append(payload),
|
||||
slack_sender=lambda payload: calls["slack"].append(payload),
|
||||
)
|
||||
result = service.notify("completed", "task-1", "proj-1", "done")
|
||||
self.assertEqual(result, {"sent": True, "channels": ["webhook", "slack"]})
|
||||
self.assertEqual(len(calls["webhook"]), 1)
|
||||
self.assertEqual(len(calls["slack"]), 1)
|
||||
self.assertEqual(calls["webhook"][0]["event_type"], "completed")
|
||||
self.assertEqual(calls["webhook"][0]["task_id"], "task-1")
|
||||
|
||||
def test_disallowed_event_not_sent_and_senders_not_called(self):
|
||||
calls = []
|
||||
service = NotificationService(
|
||||
webhook_sender=lambda payload: calls.append(("webhook", payload)),
|
||||
slack_sender=lambda payload: calls.append(("slack", payload)),
|
||||
)
|
||||
result = service.notify("running", "task-1", "proj-1")
|
||||
self.assertEqual(result, {"sent": False, "reason": "event not allowed"})
|
||||
self.assertEqual(calls, [])
|
||||
|
||||
def test_no_senders_does_not_error(self):
|
||||
service = NotificationService()
|
||||
result = service.notify("completed", "task-1", "proj-1")
|
||||
self.assertEqual(result, {"sent": True, "channels": []})
|
||||
|
||||
def test_allowed_events_set(self):
|
||||
self.assertEqual(
|
||||
ALLOWED_NOTIFICATION_EVENTS,
|
||||
{"completed", "failed", "goal-unmet", "merge-failed",
|
||||
"review-requested"},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user