- _send_json 加 extra_headers 参数,429 返回处传 Retry-After = limiter.window_seconds(默认 60),客户端据此退避一个完整窗口, 不再以原速率 hammer 把自己钉在限流里(Postgres DbRateLimiter 先 INSERT 再 COUNT,被拒请求本身也计数)。 - 单测:429 响应含 Retry-After 且值等于窗口秒数(默认 60 与注入 窗口 7 各验一次)。 - 相关单测:test_prod_webhook_exec + test_rate_limit_db 21 个、 test_server + test_server_endpoints 28 个,全过;cloud 全量未跑。
334 lines
14 KiB
Python
334 lines
14 KiB
Python
"""生产级 webhook / 备份 / 限流 / 指标单测。
|
||
|
||
只用标准库,不跑真实模型回合、不联网、不连真实 GitHub;
|
||
webhook secret 均为单测内临时值,不是生产密钥。
|
||
"""
|
||
|
||
import hashlib
|
||
import hmac
|
||
import http.client
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import unittest
|
||
from pathlib import Path
|
||
from unittest import mock
|
||
|
||
try:
|
||
from cloud.auth import AuthStore
|
||
from cloud.server import RateLimiter, create_server
|
||
from cloud.webhooks import WebhookDeliveryStore, parse_github_check_run
|
||
except ImportError:
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.auth import AuthStore
|
||
from cloud.server import RateLimiter, create_server
|
||
from cloud.webhooks import WebhookDeliveryStore, parse_github_check_run
|
||
|
||
|
||
def _sig(secret, raw):
|
||
return "sha256=" + hmac.new(
|
||
secret.encode("utf-8"), raw, hashlib.sha256).hexdigest()
|
||
|
||
|
||
class WebhookExecTestCase(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._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
|
||
os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
|
||
# P1:备份 out_dir 限白名单根;本测试的 tmp 目录即白名单根。
|
||
self._old_backup_root = os.environ.get("CLOUD_BACKUP_ROOT")
|
||
os.environ["CLOUD_BACKUP_ROOT"] = self.tmp.name
|
||
self.viewer_id, self.viewer_key = self.auth.create_user("viewer")
|
||
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()
|
||
# webhook 任务归属按项目名 == repo full_name 匹配。
|
||
status, self.project = self._req(
|
||
"POST", "/projects", {"name": "owner/repo"}, self.owner_key)
|
||
self.assertEqual(status, 201, self.project)
|
||
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()
|
||
if getattr(self, "_old_admins", None) is None:
|
||
os.environ.pop("CODEX_GLOBAL_ADMINS", None)
|
||
else:
|
||
os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins
|
||
if getattr(self, "_old_backup_root", None) is None:
|
||
os.environ.pop("CLOUD_BACKUP_ROOT", None)
|
||
else:
|
||
os.environ["CLOUD_BACKUP_ROOT"] = self._old_backup_root
|
||
|
||
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 _webhook(self, payload, delivery="d-1", event="issue_comment",
|
||
secret="test-secret"):
|
||
raw = json.dumps(payload).encode("utf-8")
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"X-Hub-Signature-256": _sig(secret, raw),
|
||
"X-GitHub-Event": event,
|
||
"X-GitHub-Delivery": delivery,
|
||
}
|
||
with mock.patch.dict(
|
||
os.environ, {"GITHUB_WEBHOOK_SECRET": secret}):
|
||
return self._req(
|
||
"POST", "/webhooks/github", raw_body=raw, headers=headers)
|
||
|
||
def _fix_payload(self):
|
||
return {
|
||
"action": "created",
|
||
"comment": {"body": "请 /fix-review 看一下",
|
||
"user": {"login": "alice",
|
||
"author_association": "COLLABORATOR"}},
|
||
"repository": {"full_name": "owner/repo"},
|
||
"issue": {"number": 42, "user": {"login": "alice"}},
|
||
}
|
||
|
||
def _tasks(self):
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
return conn.execute(
|
||
"SELECT id, prompt, status FROM tasks").fetchall()
|
||
finally:
|
||
conn.close()
|
||
|
||
def _audit_actions(self):
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
return [r[0] for r in conn.execute(
|
||
"SELECT action FROM audit_log").fetchall()]
|
||
finally:
|
||
conn.close()
|
||
|
||
def test_fix_review_creates_queued_task_and_audit(self):
|
||
status, body = self._webhook(self._fix_payload(), delivery="d-100")
|
||
self.assertEqual(status, 200, body)
|
||
self.assertEqual(body["type"], "fix_review")
|
||
self.assertIsNotNone(body["task_id"])
|
||
self.assertEqual(body["project_id"], self.project["id"])
|
||
tasks = self._tasks()
|
||
self.assertEqual(len(tasks), 1)
|
||
task_id, prompt, task_status = tasks[0]
|
||
self.assertEqual(task_id, body["task_id"])
|
||
self.assertEqual(task_status, "queued")
|
||
self.assertIn("owner/repo", prompt)
|
||
self.assertIn("42", prompt)
|
||
self.assertIn("请 /fix-review 看一下", prompt)
|
||
self.assertIn("alice", prompt)
|
||
self.assertIn("webhook_fix_review", self._audit_actions())
|
||
|
||
def test_duplicate_delivery_creates_no_second_task_or_audit(self):
|
||
status, first = self._webhook(self._fix_payload(), delivery="d-dup")
|
||
self.assertEqual(status, 200)
|
||
self.assertIsNotNone(first["task_id"])
|
||
status, second = self._webhook(self._fix_payload(), delivery="d-dup")
|
||
self.assertEqual(status, 200)
|
||
self.assertTrue(second["duplicate"])
|
||
self.assertEqual(len(self._tasks()), 1)
|
||
self.assertEqual(
|
||
self._audit_actions().count("webhook_fix_review"), 1)
|
||
|
||
def test_fix_review_no_matching_project_creates_no_task(self):
|
||
payload = self._fix_payload()
|
||
payload["repository"] = {"full_name": "other/repo"}
|
||
status, body = self._webhook(payload, delivery="d-nomatch")
|
||
self.assertEqual(status, 200)
|
||
self.assertIsNone(body["task_id"])
|
||
self.assertEqual(self._tasks(), [])
|
||
|
||
def test_check_run_success_and_failure_parsed_no_task(self):
|
||
for conclusion in ("success", "failure"):
|
||
payload = {
|
||
"repository": {"full_name": "owner/repo"},
|
||
"check_run": {"name": "verify", "head_sha": "abc123",
|
||
"conclusion": conclusion},
|
||
}
|
||
status, body = self._webhook(
|
||
payload, delivery="d-check-%s" % conclusion,
|
||
event="check_run")
|
||
self.assertEqual(status, 200, body)
|
||
self.assertEqual(body["type"], "check_run")
|
||
self.assertEqual(body["repo"], "owner/repo")
|
||
self.assertEqual(body["head_sha"], "abc123")
|
||
self.assertEqual(body["conclusion"], conclusion)
|
||
self.assertEqual(body["check_name"], "verify")
|
||
self.assertEqual(self._tasks(), [])
|
||
|
||
def test_check_run_in_progress_returns_no_check_result(self):
|
||
payload = {
|
||
"repository": {"full_name": "owner/repo"},
|
||
"check_run": {"name": "verify", "head_sha": "abc123",
|
||
"conclusion": None},
|
||
}
|
||
status, body = self._webhook(
|
||
payload, delivery="d-check-prog", event="check_run")
|
||
self.assertEqual(status, 200)
|
||
self.assertEqual(body, {"result": None})
|
||
|
||
def test_webhook_bad_signature_401_and_secret_missing_503(self):
|
||
payload = self._fix_payload()
|
||
raw = json.dumps(payload).encode("utf-8")
|
||
with mock.patch.dict(
|
||
os.environ, {"GITHUB_WEBHOOK_SECRET": "test-secret"}):
|
||
status, _ = self._req(
|
||
"POST", "/webhooks/github", raw_body=raw,
|
||
headers={"Content-Type": "application/json",
|
||
"X-Hub-Signature-256": "sha256=deadbeef",
|
||
"X-GitHub-Event": "issue_comment",
|
||
"X-GitHub-Delivery": "d-bad"})
|
||
self.assertEqual(status, 401)
|
||
with mock.patch.dict(os.environ, {}, clear=False):
|
||
os.environ.pop("GITHUB_WEBHOOK_SECRET", None)
|
||
status, _ = self._req(
|
||
"POST", "/webhooks/github", raw_body=raw,
|
||
headers={"Content-Type": "application/json",
|
||
"X-Hub-Signature-256": _sig("test-secret", raw)})
|
||
self.assertEqual(status, 503)
|
||
|
||
def test_admin_backup_sqlite_owner_ok_viewer_forbidden(self):
|
||
out_dir = os.path.join(self.tmp.name, "backups")
|
||
status, body = self._req(
|
||
"POST", "/admin/backup", {"out_dir": out_dir}, self.owner_key)
|
||
self.assertEqual(status, 200, body)
|
||
self.assertEqual(body["backend"], "sqlite")
|
||
self.assertTrue(os.path.exists(body["path"]))
|
||
self.assertGreater(body["size_bytes"], 0)
|
||
self.assertIn("task_count", body)
|
||
self.assertNotIn("password", json.dumps(body))
|
||
status, _ = self._req(
|
||
"POST", "/admin/backup", {"out_dir": out_dir}, self.viewer_key)
|
||
self.assertEqual(status, 403)
|
||
status, _ = self._req(
|
||
"POST", "/admin/backup", {"out_dir": out_dir}, key=None)
|
||
self.assertEqual(status, 401)
|
||
|
||
def test_rate_limit_429_not_401(self):
|
||
# 新限流器 + 小阈值注入:前 3 次通,第 4 次 429。
|
||
self.server.rate_limiter = RateLimiter(per_minute=3)
|
||
self.server.rate_limit_per_minute = 3
|
||
for _ in range(3):
|
||
status, _ = self._req(
|
||
"GET", "/projects", key=self.owner_key)
|
||
self.assertEqual(status, 200)
|
||
status, body = self._req("GET", "/projects", key=self.owner_key)
|
||
self.assertEqual(status, 429)
|
||
self.assertEqual(body, {"error": "rate_limited"})
|
||
|
||
def test_rate_limit_429_has_retry_after_window_seconds(self):
|
||
# 429 响应必须含 Retry-After 头,且值为限流窗口秒数
|
||
# (limiter.window_seconds),不是硬编码 60:先用默认
|
||
# 窗口 60 验一次,再用注入窗口 7 验一次。
|
||
for window, per_minute in ((60, 1), (7, 1)):
|
||
self.server.rate_limiter = RateLimiter(
|
||
per_minute=per_minute, window_seconds=window)
|
||
self.server.rate_limit_per_minute = per_minute
|
||
self.server.rate_limit_window_seconds = window
|
||
conn = http.client.HTTPConnection(
|
||
"127.0.0.1", self.port, timeout=10)
|
||
try:
|
||
conn.request("GET", "/projects",
|
||
headers={"X-API-Key": self.owner_key})
|
||
resp = conn.getresponse()
|
||
resp.read()
|
||
self.assertEqual(resp.status, 200)
|
||
conn.request("GET", "/projects",
|
||
headers={"X-API-Key": self.owner_key})
|
||
resp = conn.getresponse()
|
||
body = resp.read()
|
||
self.assertEqual(resp.status, 429)
|
||
self.assertEqual(json.loads(body.decode("utf-8")),
|
||
{"error": "rate_limited"})
|
||
self.assertEqual(resp.getheader("Retry-After"),
|
||
str(window))
|
||
finally:
|
||
conn.close()
|
||
|
||
def test_health_backend_and_db_ok_and_metrics_by_status(self):
|
||
status, body = self._req("GET", "/health")
|
||
self.assertEqual(status, 200)
|
||
self.assertTrue(body["ok"])
|
||
self.assertEqual(body["backend"], "sqlite")
|
||
self.assertTrue(body["db_ok"])
|
||
status, metrics = self._req(
|
||
"GET", "/metrics", key=self.owner_key)
|
||
self.assertEqual(status, 200)
|
||
for key in ("queued", "running", "completed", "failed",
|
||
"interrupted", "merged"):
|
||
self.assertIn(key, metrics["by_status"])
|
||
|
||
|
||
class WebhookUnitTestCase(unittest.TestCase):
|
||
def test_parse_check_run_only_success_failure(self):
|
||
base = {"repository": {"full_name": "o/r"}}
|
||
ok = dict(base, check_run={"name": "v", "head_sha": "h",
|
||
"conclusion": "success"})
|
||
self.assertEqual(parse_github_check_run(ok)["conclusion"], "success")
|
||
bad = dict(base, check_run={"name": "v", "head_sha": "h",
|
||
"conclusion": "neutral"})
|
||
self.assertIsNone(parse_github_check_run(bad))
|
||
self.assertIsNone(parse_github_check_run({}))
|
||
suite = dict(base, check_suite={"head_sha": "h2",
|
||
"conclusion": "failure"})
|
||
parsed = parse_github_check_run(suite)
|
||
self.assertEqual(parsed["conclusion"], "failure")
|
||
self.assertEqual(parsed["head_sha"], "h2")
|
||
|
||
def test_delivery_store_claim_duplicate(self):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
db_path = os.path.join(tmp.name, "c.db")
|
||
store = WebhookDeliveryStore(db_path)
|
||
try:
|
||
self.assertTrue(store.claim("d1", event="issue_comment",
|
||
repo="o/r"))
|
||
self.assertFalse(store.claim("d1", event="issue_comment",
|
||
repo="o/r"))
|
||
self.assertTrue(store.claim("d2"))
|
||
finally:
|
||
store.close()
|
||
tmp.cleanup()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|