host
4ed1479d7a
- select_final trivial condition now k==1 and single-candidate pool, regardless of whether a selector is provided (previously gated on selector is None, so wired runs always paid for a useless selector task at k=1) - Budget gains max_sessions gate (probe + stage-attempt sessions; fmt2 corrections and selector/reflector sessions exempt), the backstop for unpriced models where the USD gate never binds - validate_run_tag (charset, no leading '-', no s<N> seed prefix) enforced at run_improvement_via_tasks and the CLI - CLI --check preflight (base harness, source dir, prompt, tag, control-plane /health probe) without creating tasks - ModeResult.to_dict now includes rollout summaries so a lost probe is visible in full_result instead of only on the server - CLI plan counts the selector only when k>1
568 lines
25 KiB
Python
568 lines
25 KiB
Python
"""cloud/harness2_tasks.py 单测：真实 codex 任务版接线。
|
||
|
||
不跑真实模型、不联网：FakeTaskClient 在进程内模拟控制面——探针
|
||
任务真跑 bundle 里的脚本，编辑任务按文件块格式回改动，选择与
|
||
反思按最终回复格式回结果。HTTP 客户端则对本地假服务端实测。
|
||
账本与 base 树同 test_harness2.py：带符号和 85.15，只加正数 105.25。
|
||
"""
|
||
|
||
import csv
|
||
import json
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import unittest
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
from pathlib import Path
|
||
|
||
try:
|
||
from cloud import harness2_tasks
|
||
from cloud.harness2 import (
|
||
parse_choice, parse_experience, parse_harness_files, parse_no_change,
|
||
parse_notes, parse_why,
|
||
)
|
||
from cloud.harness2_tasks import (
|
||
CodexEditor, CodexProbeAdapter, HttpTaskClient, TaskState,
|
||
build_multipart, preflight_checks, run_improvement_via_tasks, run_task,
|
||
unwrap_result_text, validate_run_tag,
|
||
)
|
||
from cloud.uploads import parse_multipart, parse_paths_field
|
||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud import harness2_tasks
|
||
from cloud.harness2 import (
|
||
parse_choice, parse_experience, parse_harness_files, parse_no_change,
|
||
parse_notes, parse_why,
|
||
)
|
||
from cloud.harness2_tasks import (
|
||
CodexEditor, CodexProbeAdapter, HttpTaskClient, TaskState,
|
||
build_multipart, preflight_checks, run_improvement_via_tasks, run_task,
|
||
unwrap_result_text, validate_run_tag,
|
||
)
|
||
from cloud.uploads import parse_multipart, parse_paths_field
|
||
|
||
BASE_SCRIPT = """import csv
|
||
import sys
|
||
from decimal import Decimal
|
||
with open(sys.argv[1]) as f:
|
||
amounts = [Decimal(r['amount']) for r in csv.DictReader(f)]
|
||
print(f"Total: {sum((a for a in amounts if a > 0), Decimal(0)):.2f}")
|
||
"""
|
||
|
||
BASE_TREE = {
|
||
"AGENTS.md": "Use scripts/summarize.py to total the ledger.\n",
|
||
"scripts/summarize.py": BASE_SCRIPT,
|
||
}
|
||
|
||
LEDGER = b"amount\n100.25\n-20.10\n5.00\n"
|
||
|
||
|
||
class FakeTaskClient:
|
||
"""进程内控制面：按 prompt 类型分派到确定性假 worker。"""
|
||
|
||
def __init__(self, editor_mode="fix", selector_mode="signed", envelope=False):
|
||
self.editor_mode = editor_mode
|
||
self.selector_mode = selector_mode
|
||
self.envelope = envelope
|
||
self.uploads = {}
|
||
self.tasks = {}
|
||
self.by_key = {}
|
||
self.created = 0
|
||
self.cancelled = []
|
||
|
||
def upload(self, files):
|
||
uid = f"u{len(self.uploads) + 1}"
|
||
self.uploads[uid] = dict(files)
|
||
return uid
|
||
|
||
def create_task(self, *, project_id, prompt, upload_id=None,
|
||
idempotency_key=None, **kwargs):
|
||
if idempotency_key and idempotency_key in self.by_key:
|
||
return self.by_key[idempotency_key]
|
||
tid = f"t{len(self.tasks) + 1}"
|
||
self.created += 1
|
||
files = self.uploads.get(upload_id, {}) if upload_id else {}
|
||
result = self._work(prompt, files)
|
||
if self.envelope:
|
||
# 模拟控制面真实形状：result_text 是 JSON 信封，text 字段
|
||
# 里才是最终回复——接线层必须拆信封，否则行首锚定的文件
|
||
# 块解析永远看不见块（2026-10-29 真实轮次根因）。
|
||
result = json.dumps({
|
||
"task": tid, "thread_id": "th", "cwd": "/x",
|
||
"status": "completed", "text": result, "error": None,
|
||
})
|
||
self.tasks[tid] = TaskState(task_id=tid, status="completed",
|
||
result_text=result)
|
||
if idempotency_key:
|
||
self.by_key[idempotency_key] = tid
|
||
return tid
|
||
|
||
def _work(self, prompt, files):
|
||
if prompt.startswith("You are editing an agent harness"):
|
||
if self.editor_mode == "gibberish":
|
||
return "hello, I did stuff"
|
||
if self.editor_mode in ("notes-only-then-fix", "notes-only"):
|
||
if ("YOUR PREVIOUS REPLY WAS REJECTED" in prompt
|
||
and self.editor_mode == "notes-only-then-fix"):
|
||
script = files["context/harness/scripts/summarize.py"].decode()
|
||
fixed = script.replace(" if a > 0", "")
|
||
return (
|
||
"=== HARNESS FILE: scripts/summarize.py ===\n" + fixed +
|
||
"=== END HARNESS FILE ===\n"
|
||
"<notes>Removed the positive-only filter.</notes>")
|
||
# 真实轮次原样复现：notes 里声称已修正，却不发文件块。
|
||
return ("<notes>Removed the positive-only filter so signed "
|
||
"amounts are summed; the script now totals all amounts."
|
||
"</notes>")
|
||
if self.editor_mode == "no-change":
|
||
return ("<notes>No evidence justifies a change.</notes>"
|
||
"<no-change>The base trajectory shows no defect."
|
||
"</no-change>")
|
||
script = files["context/harness/scripts/summarize.py"].decode()
|
||
fixed = script.replace(" if a > 0", "")
|
||
return (
|
||
"=== HARNESS FILE: scripts/summarize.py ===\n" + fixed +
|
||
"=== END HARNESS FILE ===\n"
|
||
"<notes>Keep signed amounts: the base total omitted credits.</notes>"
|
||
)
|
||
if prompt.startswith("You are choosing which"):
|
||
if self.selector_mode == "gibberish":
|
||
return "I like them all"
|
||
cands = sorted(p for p in files if p.startswith("context/candidates/")
|
||
and p.endswith(".md"))
|
||
for p in cands:
|
||
if b"Total: 85.15" in files[p]:
|
||
stem = Path(p).stem
|
||
return (f"```choice\n{stem}\n```\n"
|
||
"<why>The signed total matches the source ledger.</why>")
|
||
stem = Path(cands[0]).stem
|
||
return (f"```choice\n{stem}\n```\n<why>first delivered candidate</why>")
|
||
if prompt.startswith("You are reading the complete record"):
|
||
changed = (files.get("context/harness/scripts/summarize.py")
|
||
!= files.get("context/base_harness/scripts/summarize.py"))
|
||
if not changed:
|
||
return "NO_ENTRY"
|
||
return (
|
||
"<edits>\nscripts/summarize.py: dropped the positive-only filter; "
|
||
"verdict: helped\n</edits>\n"
|
||
"<experience>\ncue: When a ledger includes credits and debits\n"
|
||
"lesson: Include both signs when summing in scripts/summarize.py.\n"
|
||
"evidence: Removing the positive-only filter changed the probe total "
|
||
"to include the credits omitted by the base trajectory.\n"
|
||
"confidence: high\n</experience>"
|
||
)
|
||
# 探针：真跑 bundle 里的脚本对 bundle 里的账本求和。
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
root = Path(tmp)
|
||
for rel, data in files.items():
|
||
dest = root / rel
|
||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||
dest.write_bytes(data)
|
||
proc = subprocess.run(
|
||
[sys.executable, str(root / "scripts" / "summarize.py"),
|
||
str(root / "ledger.csv")],
|
||
capture_output=True, text=True, timeout=10)
|
||
return proc.stdout.strip()
|
||
|
||
def get_task(self, task_id):
|
||
return self.tasks[task_id]
|
||
|
||
def get_trace_text(self, task_id):
|
||
return "[ACTION] fake probe\n[OBSERVATION]\n" + self.tasks[task_id].result_text
|
||
|
||
def get_cost_usd(self, task_id):
|
||
return 0.02
|
||
|
||
def cancel_task(self, task_id):
|
||
self.cancelled.append(task_id)
|
||
|
||
|
||
class TestParseHelpers(unittest.TestCase):
|
||
def test_multipart_roundtrips_through_server_parser(self):
|
||
files = {"context/task.md": b"do it", "AGENTS.md": b"be nice\n"}
|
||
body, content_type = build_multipart(files)
|
||
fields, parsed = parse_multipart(body, content_type)
|
||
paths = parse_paths_field(fields, parsed)
|
||
self.assertEqual(paths, sorted(files))
|
||
self.assertEqual(dict(parsed), files)
|
||
|
||
def test_choice_why_notes_experience(self):
|
||
self.assertEqual(parse_choice("```choice\nB\n```"), "B")
|
||
self.assertIsNone(parse_choice("no choice here"))
|
||
self.assertEqual(parse_why("<why>because</why>"), "because")
|
||
self.assertEqual(parse_notes("<notes>did x</notes>"), "did x")
|
||
exp = parse_experience(
|
||
"<experience>\ncue: c\nlesson: l\nevidence: e\nconfidence: high\n</experience>")
|
||
self.assertEqual((exp["cue"], exp["lesson"], exp["evidence"]), ("c", "l", "e"))
|
||
self.assertEqual(parse_experience("NO_ENTRY"), {})
|
||
self.assertEqual(parse_experience("<experience>\ncue: c\n</experience>"), {})
|
||
|
||
def test_harness_file_blocks(self):
|
||
text = ("=== HARNESS FILE: AGENTS.md ===\nhello\n"
|
||
"=== END HARNESS FILE ===\n<notes>n</notes>")
|
||
self.assertEqual(parse_harness_files(text), {"AGENTS.md": "hello\n"})
|
||
self.assertEqual(parse_harness_files("nothing here"), {})
|
||
with self.assertRaises(ValueError):
|
||
parse_harness_files("=== HARNESS FILE: /etc/passwd ===\nx\n"
|
||
"=== END HARNESS FILE ===")
|
||
with self.assertRaises(ValueError):
|
||
parse_harness_files("=== HARNESS FILE: AGENTS.md ===\nunterminated\n")
|
||
|
||
|
||
class TestFullRunViaTasks(unittest.TestCase):
|
||
def _run(self, client, run_tag="run1", work_sub="work", k=1, max_sessions=None):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
self.addCleanup(tmp.cleanup)
|
||
root = Path(tmp.name)
|
||
result = run_improvement_via_tasks(
|
||
client, project_id="p1", task_id="ledger-1",
|
||
prompt="Sum every signed amount in ledger.csv. Return Total: with two decimals.",
|
||
base_harness=dict(BASE_TREE), run_tag=run_tag,
|
||
work_dir=root / work_sub, domain="accounting",
|
||
source_files={"ledger.csv": LEDGER},
|
||
k=k, mode="parallel", library_root=root / "library",
|
||
max_sessions=max_sessions)
|
||
return result, client, root
|
||
|
||
def test_parallel_k1_via_real_task_shape(self):
|
||
client = FakeTaskClient()
|
||
result, _, root = self._run(client)
|
||
self.assertTrue(result.ok)
|
||
self.assertEqual(result.selection.chosen_label, "agg1")
|
||
self.assertTrue(result.selection.trivial) # k=1 单终版不开选择会话
|
||
self.assertEqual(result.mode_result.rollouts["agg1"].deliverable_text,
|
||
"Total: 85.15")
|
||
# 4 探针 + 3 编辑 + 1 反思 = 8 个任务（选择 trivial，不建任务）
|
||
self.assertEqual(client.created, 8)
|
||
self.assertIsNotNone(result.library_entry)
|
||
self.assertTrue((root / "library" / "accounting" / "INDEX.md").is_file())
|
||
|
||
def test_same_run_tag_replays_without_new_tasks(self):
|
||
client = FakeTaskClient()
|
||
result1, _, _ = self._run(client, work_sub="work1")
|
||
self.assertTrue(result1.ok)
|
||
created_after_first = client.created
|
||
result2, _, _ = self._run(client, work_sub="work2")
|
||
self.assertTrue(result2.ok)
|
||
self.assertEqual(client.created, created_after_first) # 幂等键命中，零新建
|
||
|
||
def test_editor_gibberish_falls_back(self):
|
||
client = FakeTaskClient(editor_mode="gibberish")
|
||
result, _, _ = self._run(client)
|
||
self.assertTrue(result.ok) # 回落的 base 产物仍可交付
|
||
self.assertEqual(set(result.mode_result.fell_back), {"inv1", "dom1", "agg1"})
|
||
self.assertIsNone(result.library_entry)
|
||
|
||
def test_selector_gibberish_falls_back(self):
|
||
# k=1 时选择 trivial，不走选择器；用 k=2 让选择器真正跑起来。
|
||
client = FakeTaskClient(selector_mode="gibberish")
|
||
result, _, _ = self._run(client, k=2)
|
||
self.assertTrue(result.ok)
|
||
self.assertTrue(result.selection.fallback)
|
||
|
||
|
||
class TestEditorFormatRetry(unittest.TestCase):
|
||
"""编辑格式纠正：notes-only 先重试再落盘；显式无改动一次通过；
|
||
两次都 notes-only 则如实抛错（阶段回落，不静默采纳）。"""
|
||
|
||
def _editor(self, mode):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
self.addCleanup(tmp.cleanup)
|
||
ctx = Path(tmp.name)
|
||
(ctx / "harness" / "scripts").mkdir(parents=True)
|
||
for rel, content in BASE_TREE.items():
|
||
(ctx / "harness" / rel).write_text(content)
|
||
client = FakeTaskClient(editor_mode=mode)
|
||
editor = CodexEditor(client, project_id="p", run_tag="fmt")
|
||
return editor, ctx, client
|
||
|
||
def test_notes_only_then_fix_lands(self):
|
||
editor, ctx, client = self._editor("notes-only-then-fix")
|
||
result = editor.edit("You are editing an agent harness", ctx, "inv1.a1")
|
||
self.assertNotIn("if a > 0", result.tree["scripts/summarize.py"])
|
||
self.assertEqual(client.created, 2) # 首发 + 一次格式纠正
|
||
|
||
def test_no_change_marker_keeps_base(self):
|
||
editor, ctx, client = self._editor("no-change")
|
||
result = editor.edit("You are editing an agent harness", ctx, "inv1.a1")
|
||
self.assertEqual(result.tree["scripts/summarize.py"], BASE_SCRIPT)
|
||
self.assertEqual(client.created, 1)
|
||
|
||
def test_notes_only_twice_raises(self):
|
||
editor, ctx, client = self._editor("notes-only")
|
||
with self.assertRaises(ValueError):
|
||
editor.edit("You are editing an agent harness", ctx, "inv1.a1")
|
||
self.assertEqual(client.created, 2)
|
||
|
||
def test_parse_no_change(self):
|
||
self.assertEqual(parse_no_change("<no-change>no defect</no-change>"),
|
||
"no defect")
|
||
self.assertEqual(parse_no_change("just notes"), "")
|
||
|
||
|
||
class TestUnwrapResultText(unittest.TestCase):
|
||
"""拆信封：信封取 text；其它原样返回，不丢内容。"""
|
||
|
||
def test_envelope_unwrapped(self):
|
||
raw = json.dumps({"task": "t1", "status": "completed",
|
||
"text": "=== HARNESS FILE: a.md ===\nhi\n=== END HARNESS FILE ===\n"})
|
||
self.assertTrue(unwrap_result_text(raw).startswith("=== HARNESS FILE"))
|
||
|
||
def test_plain_text_unchanged(self):
|
||
self.assertEqual(unwrap_result_text("Total: 85.15"), "Total: 85.15")
|
||
self.assertEqual(unwrap_result_text('{"no": "text field"}'),
|
||
'{"no": "text field"}')
|
||
self.assertEqual(unwrap_result_text("{not json"), "{not json")
|
||
self.assertEqual(unwrap_result_text(""), "")
|
||
|
||
|
||
class TestEnvelopeFullRound(unittest.TestCase):
|
||
"""信封形状整轮：假控制面按真实信封返回时，编辑块仍能落盘、
|
||
探针交付物是干净文本（而非信封 JSON），经验正常入库。"""
|
||
|
||
def test_envelope_round(self):
|
||
client = FakeTaskClient(envelope=True)
|
||
tmp = tempfile.TemporaryDirectory()
|
||
self.addCleanup(tmp.cleanup)
|
||
root = Path(tmp.name)
|
||
result = run_improvement_via_tasks(
|
||
client, project_id="p1", task_id="ledger-1",
|
||
prompt="Sum every signed amount in ledger.csv. Return Total: with two decimals.",
|
||
base_harness=dict(BASE_TREE), run_tag="env",
|
||
work_dir=root / "work", domain="accounting",
|
||
source_files={"ledger.csv": LEDGER},
|
||
k=1, mode="parallel", library_root=root / "library")
|
||
self.assertTrue(result.ok)
|
||
self.assertEqual(result.mode_result.rollouts["agg1"].deliverable_text,
|
||
"Total: 85.15")
|
||
self.assertNotIn("thread_id",
|
||
result.mode_result.rollouts["agg1"].deliverable_text)
|
||
self.assertIsNotNone(result.library_entry)
|
||
|
||
|
||
class TestSessionBudget(unittest.TestCase):
|
||
"""会话数门：0 挡全部、2 挡在中途；fmt2 与选择/反思不计入。"""
|
||
|
||
def _run(self, max_sessions):
|
||
client = FakeTaskClient()
|
||
tmp = tempfile.TemporaryDirectory()
|
||
self.addCleanup(tmp.cleanup)
|
||
root = Path(tmp.name)
|
||
result = run_improvement_via_tasks(
|
||
client, project_id="p1", task_id="ledger-1",
|
||
prompt="Sum every signed amount in ledger.csv. Return Total: with two decimals.",
|
||
base_harness=dict(BASE_TREE), run_tag="sess",
|
||
work_dir=root / "work", domain="accounting",
|
||
source_files={"ledger.csv": LEDGER},
|
||
k=1, mode="parallel", library_root=root / "library",
|
||
max_sessions=max_sessions)
|
||
return result, client
|
||
|
||
def test_zero_blocks_everything(self):
|
||
result, client = self._run(0)
|
||
self.assertEqual(client.created, 0)
|
||
self.assertEqual(set(result.mode_result.fell_back), {"inv1", "dom1", "agg1"})
|
||
|
||
def test_mid_cap_stops_after_two_sessions(self):
|
||
# base 探针（1）+ inv1 编辑（2），之后 dom1/agg1 与探针全被门挡住。
|
||
result, client = self._run(2)
|
||
self.assertEqual(client.created, 2)
|
||
self.assertIn("dom1", result.mode_result.fell_back)
|
||
self.assertIn("agg1", result.mode_result.fell_back)
|
||
self.assertNotIn("inv1", result.mode_result.fell_back)
|
||
|
||
|
||
class TestValidateRunTag(unittest.TestCase):
|
||
def test_valid(self):
|
||
for tag in ("h2-20261029-real1", "a", "a.b_c-1", "run_1"):
|
||
validate_run_tag(tag)
|
||
|
||
def test_invalid(self):
|
||
for tag in ("-bad", "has/slash", "has space", "s1-foo", "s12"):
|
||
with self.assertRaises(ValueError, msg=tag):
|
||
validate_run_tag(tag)
|
||
validate_run_tag("") # 空串表示稍后生成，不校验
|
||
|
||
|
||
class TestPreflight(unittest.TestCase):
|
||
def _dirs(self, agents="Use scripts/summarize.py.\n"):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
self.addCleanup(tmp.cleanup)
|
||
root = Path(tmp.name)
|
||
(root / "base").mkdir()
|
||
if agents is not None:
|
||
(root / "base" / "AGENTS.md").write_text(agents)
|
||
(root / "source").mkdir()
|
||
(root / "source" / "ledger.csv").write_bytes(LEDGER)
|
||
return root
|
||
|
||
def test_ok(self):
|
||
root = self._dirs()
|
||
errors, _ = preflight_checks(
|
||
base_harness_dir=root / "base", source_dir=root / "source",
|
||
prompt="do it", run_tag="h2-ok")
|
||
self.assertEqual(errors, [])
|
||
|
||
def test_errors(self):
|
||
root = self._dirs(agents="")
|
||
errors, _ = preflight_checks(
|
||
base_harness_dir=root / "base", source_dir=root / "missing",
|
||
prompt=" ", run_tag="s1-bad")
|
||
self.assertEqual(len(errors), 4) # AGENTS 为空、source 缺失、提示为空、tag 非法
|
||
errors, _ = preflight_checks(base_harness_dir=root / "nope", prompt="x")
|
||
self.assertTrue(any("不存在" in e for e in errors))
|
||
|
||
|
||
class TestHttpClient(unittest.TestCase):
|
||
def test_client_against_local_server(self):
|
||
seen = {"keys": [], "uploads": [], "tasks": []}
|
||
polls = {"t1": 0}
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
def log_message(self, *args):
|
||
pass
|
||
|
||
def _json(self, status, obj):
|
||
raw = json.dumps(obj).encode()
|
||
self.send_response(status)
|
||
self.send_header("Content-Type", "application/json")
|
||
self.send_header("Content-Length", str(len(raw)))
|
||
self.end_headers()
|
||
self.wfile.write(raw)
|
||
|
||
def do_POST(self):
|
||
seen["keys"].append(self.headers.get("X-API-Key"))
|
||
length = int(self.headers.get("Content-Length") or 0)
|
||
body = self.rfile.read(length) if length else b""
|
||
if self.path == "/uploads":
|
||
seen["uploads"].append(body)
|
||
self._json(201, {"upload_id": "u1", "files": []})
|
||
elif self.path == "/tasks":
|
||
seen["tasks"].append(json.loads(body.decode()))
|
||
self._json(201, {"id": "t1"})
|
||
elif self.path == "/tasks/t1/cancel":
|
||
self._json(200, {})
|
||
else:
|
||
self._json(404, {"error": "not found"})
|
||
|
||
def do_GET(self):
|
||
if self.path == "/tasks/t1":
|
||
polls["t1"] += 1
|
||
status = "queued" if polls["t1"] == 1 else "completed"
|
||
self._json(200, {"id": "t1", "status": status,
|
||
"result_text": "Total: 85.15" if status == "completed" else "",
|
||
"error_text": ""})
|
||
elif self.path == "/tasks/t1/cost":
|
||
self._json(200, {"task_id": "t1", "usd_micros": 2500000})
|
||
elif self.path == "/tasks/t1/trace":
|
||
self._json(200, {"entries": [{"id": "e1", "kind": "agent",
|
||
"text": "hello"}]})
|
||
else:
|
||
self._json(404, {"error": "not found"})
|
||
|
||
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||
thread.start()
|
||
try:
|
||
client = HttpTaskClient(f"http://127.0.0.1:{server.server_address[1]}",
|
||
"test-key")
|
||
done = run_task(client, project_id="p1", prompt="do it",
|
||
files={"a.txt": b"hi"}, poll_seconds=0.01,
|
||
timeout_seconds=10)
|
||
finally:
|
||
server.shutdown()
|
||
server.server_close()
|
||
self.assertEqual(done.state.status, "completed")
|
||
self.assertEqual(done.state.result_text, "Total: 85.15")
|
||
self.assertEqual(done.cost_usd, 2.5)
|
||
self.assertIn("hello", done.trace_text)
|
||
self.assertTrue(all(k == "test-key" for k in seen["keys"]))
|
||
self.assertEqual(seen["tasks"][0]["project_id"], "p1")
|
||
fields, files = parse_multipart(seen["uploads"][0],
|
||
"multipart/form-data; boundary=" +
|
||
seen["uploads"][0].split(b"\r\n")[0][2:].decode())
|
||
self.assertEqual(dict(files), {"a.txt": b"hi"})
|
||
|
||
|
||
class TestHttpClientRetry(unittest.TestCase):
|
||
"""瞬断重试：掉连接与 503 各重试一次后成功；401 只请求一次。"""
|
||
|
||
def _serve(self, handler_cls):
|
||
server = ThreadingHTTPServer(("127.0.0.1", 0), handler_cls)
|
||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||
thread.start()
|
||
self.addCleanup(server.shutdown)
|
||
self.addCleanup(server.server_close)
|
||
return HttpTaskClient(
|
||
f"http://127.0.0.1:{server.server_address[1]}", "test-key",
|
||
timeout=5)
|
||
|
||
def test_drop_then_success(self):
|
||
counts = {"n": 0}
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
def log_message(self, *args):
|
||
pass
|
||
|
||
def do_GET(self):
|
||
counts["n"] += 1
|
||
if counts["n"] == 1:
|
||
self.close_connection = True
|
||
return
|
||
raw = json.dumps({"ok": True}).encode()
|
||
self.send_response(200)
|
||
self.send_header("Content-Type", "application/json")
|
||
self.send_header("Content-Length", str(len(raw)))
|
||
self.end_headers()
|
||
self.wfile.write(raw)
|
||
|
||
client = self._serve(Handler)
|
||
self.assertEqual(client._request("GET", "/flaky"), {"ok": True})
|
||
self.assertEqual(counts["n"], 2)
|
||
|
||
def test_503_then_success(self):
|
||
counts = {"n": 0}
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
def log_message(self, *args):
|
||
pass
|
||
|
||
def do_GET(self):
|
||
counts["n"] += 1
|
||
raw = json.dumps({"ok": counts["n"] > 1}).encode()
|
||
self.send_response(503 if counts["n"] == 1 else 200)
|
||
self.send_header("Content-Type", "application/json")
|
||
self.send_header("Content-Length", str(len(raw)))
|
||
self.end_headers()
|
||
self.wfile.write(raw)
|
||
|
||
client = self._serve(Handler)
|
||
self.assertEqual(client._request("GET", "/flaky"), {"ok": True})
|
||
self.assertEqual(counts["n"], 2)
|
||
|
||
def test_401_not_retried(self):
|
||
counts = {"n": 0}
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
def log_message(self, *args):
|
||
pass
|
||
|
||
def do_GET(self):
|
||
counts["n"] += 1
|
||
raw = json.dumps({"error": "unauthorized"}).encode()
|
||
self.send_response(401)
|
||
self.send_header("Content-Type", "application/json")
|
||
self.send_header("Content-Length", str(len(raw)))
|
||
self.end_headers()
|
||
self.wfile.write(raw)
|
||
|
||
client = self._serve(Handler)
|
||
with self.assertRaises(RuntimeError):
|
||
client._request("GET", "/flaky")
|
||
self.assertEqual(counts["n"], 1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|