143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
from __future__ import annotations
|
||
|
||
import contextlib
|
||
import io
|
||
import json
|
||
import sys
|
||
import unittest
|
||
from pathlib import Path
|
||
from unittest import mock
|
||
|
||
|
||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||
ACK_SCRIPTS = REPO_ROOT / "skills" / "ack" / "scripts"
|
||
sys.path.insert(0, str(ACK_SCRIPTS))
|
||
|
||
import launch_worker # noqa: E402
|
||
import worker_probe # noqa: E402
|
||
|
||
|
||
def dispatch_show(heartbeat: str | None = None) -> dict:
|
||
dispatch: dict = {
|
||
"id": "ctx_dispatch_1",
|
||
"status": "dispatched",
|
||
"last_heartbeat_at": heartbeat,
|
||
}
|
||
return {"ok": True, "result": {"dispatch": dispatch}}
|
||
|
||
|
||
def terminal_read(tail: list[str]) -> dict:
|
||
return {
|
||
"ok": True,
|
||
"result": {
|
||
"terminal": {
|
||
"handle": "term_worker_1",
|
||
"status": "running",
|
||
"tail": tail,
|
||
}
|
||
},
|
||
}
|
||
|
||
|
||
class WorkerProbeClassifyTests(unittest.TestCase):
|
||
def test_approval_stall_is_detected(self) -> None:
|
||
result = worker_probe.classify(
|
||
"Switch to gpt-5.6-luna for lower credit usage? › 1. Switch",
|
||
None,
|
||
)
|
||
self.assertEqual(result["status"], "stall")
|
||
self.assertEqual(result["stallReason"], "model-switch")
|
||
|
||
def test_approve_prompt_is_detected(self) -> None:
|
||
result = worker_probe.classify(
|
||
"Allow tool: bash\nRun this command? (y/n)", None
|
||
)
|
||
self.assertEqual(result["status"], "stall")
|
||
self.assertEqual(result["stallReason"], "approval")
|
||
|
||
def test_usage_limit_is_detected(self) -> None:
|
||
result = worker_probe.classify(
|
||
"You've hit your usage limit. Upgrade to Plus to continue.", None
|
||
)
|
||
self.assertEqual(result["status"], "stall")
|
||
self.assertEqual(result["stallReason"], "usage-limit")
|
||
|
||
def test_heartbeat_means_progress_even_without_tail_markers(self) -> None:
|
||
result = worker_probe.classify(
|
||
"some unclassified output", "2026-08-23T12:00:00Z"
|
||
)
|
||
self.assertEqual(result["status"], "progress")
|
||
self.assertEqual(result["heartbeatAt"], "2026-08-23T12:00:00Z")
|
||
|
||
def test_working_marker_means_running_without_heartbeat(self) -> None:
|
||
result = worker_probe.classify(["• Working (12s)"], None)
|
||
self.assertEqual(result["status"], "running")
|
||
|
||
def test_welcome_screen_means_not_started(self) -> None:
|
||
result = worker_probe.classify(
|
||
["Welcome to fish, the friendly interactive shell", "Type help"],
|
||
None,
|
||
)
|
||
self.assertEqual(result["status"], "not-started")
|
||
|
||
def test_empty_tail_means_not_started(self) -> None:
|
||
result = worker_probe.classify("", None)
|
||
self.assertEqual(result["status"], "not-started")
|
||
|
||
|
||
class WorkerProbeMainTests(unittest.TestCase):
|
||
def test_main_emits_single_json_with_progress(self) -> None:
|
||
with (
|
||
mock.patch.object(
|
||
worker_probe,
|
||
"resolve_executable",
|
||
return_value=Path("/trusted/orca"),
|
||
),
|
||
mock.patch.object(
|
||
worker_probe,
|
||
"run_json",
|
||
side_effect=[
|
||
dispatch_show(heartbeat="2026-08-23T12:00:00Z"),
|
||
terminal_read(["Working (5s)"]),
|
||
],
|
||
),
|
||
):
|
||
buffer = io.StringIO()
|
||
with contextlib.redirect_stdout(buffer):
|
||
exit_code = worker_probe.main(
|
||
["--task-id", "task_1", "--terminal", "term_worker_1"]
|
||
)
|
||
payload = json.loads(buffer.getvalue())
|
||
self.assertEqual(exit_code, 0)
|
||
self.assertEqual(payload["status"], "progress")
|
||
self.assertEqual(payload["taskId"], "task_1")
|
||
self.assertEqual(payload["dispatchId"], "ctx_dispatch_1")
|
||
self.assertEqual(payload["terminal"], "term_worker_1")
|
||
|
||
def test_main_reports_unknown_on_probe_failure_without_raising(self) -> None:
|
||
with (
|
||
mock.patch.object(
|
||
worker_probe,
|
||
"resolve_executable",
|
||
return_value=Path("/trusted/orca"),
|
||
),
|
||
mock.patch.object(
|
||
worker_probe,
|
||
"run_json",
|
||
side_effect=launch_worker.LaunchError("orca unreachable"),
|
||
),
|
||
):
|
||
buffer = io.StringIO()
|
||
with contextlib.redirect_stdout(buffer):
|
||
exit_code = worker_probe.main(
|
||
["--task-id", "task_1", "--terminal", "term_worker_1"]
|
||
)
|
||
payload = json.loads(buffer.getvalue())
|
||
self.assertEqual(exit_code, 0)
|
||
self.assertEqual(payload["status"], "unknown")
|
||
self.assertIn("probe failed", payload["evidence"])
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|