wehub-resource-sync
163 lines
7.2 KiB
Python
163 lines
7.2 KiB
Python
import contextlib
|
||
import io
|
||
import json
|
||
import sys
|
||
import unittest
|
||
from unittest.mock import patch
|
||
|
||
import main_agent
|
||
|
||
|
||
class FakeSessionIndex:
|
||
def __init__(self, fail_session=False):
|
||
self.fail_session = fail_session
|
||
self.activated = ""
|
||
self.created = 0
|
||
|
||
def set_active(self, session_id):
|
||
if self.fail_session:
|
||
raise KeyError(session_id)
|
||
self.activated = session_id
|
||
|
||
def create(self):
|
||
self.created += 1
|
||
self.activated = f"new-{self.created}"
|
||
return {"session_id": self.activated}
|
||
|
||
def snapshot(self):
|
||
return {"active_session_id": self.activated, "session": {"session_id": self.activated, "stage": "created"}}
|
||
|
||
|
||
class FakeRuntime:
|
||
def __init__(self, fail_session=False):
|
||
self.session_index = FakeSessionIndex(fail_session=fail_session)
|
||
self.inputs = []
|
||
|
||
async def compact_history(self, reason="manual"):
|
||
return "Compacted context 100 -> 50 (fallback-local)."
|
||
|
||
async def stream_events(self, user_input):
|
||
self.inputs.append(user_input)
|
||
turn_id = "turn-test"
|
||
yield {"type": "turn", "turn_id": turn_id, "turn": {"id": turn_id}}
|
||
yield {"type": "status", "turn_id": turn_id, "phase": "sampling_assistant", "message": "Sampling assistant"}
|
||
yield {"type": "tool_progress", "turn_id": turn_id, "tool": {"name": "fake_tool"}, "progress": {"stage": "running", "message": "working"}}
|
||
yield {"type": "terminal", "turn_id": turn_id, "stream": "stdout", "line": "pipeline output"}
|
||
yield {"type": "token", "turn_id": turn_id, "delta": "done"}
|
||
yield {"type": "done", "turn_id": turn_id, "assistant": "done", "tool_results": []}
|
||
yield {"type": "session", "turn_id": turn_id, "session": {"session": {"session_id": "s1", "stage": "narrative_planned"}}}
|
||
|
||
|
||
class MainAgentCliTests(unittest.IsolatedAsyncioTestCase):
|
||
async def run_cli(self, argv, runtime=None, stdin_text="", session_index=None, load_runtime_side_effect=None):
|
||
runtime = runtime or FakeRuntime()
|
||
session_index = session_index or runtime.session_index
|
||
stdout = io.StringIO()
|
||
stderr = io.StringIO()
|
||
stdin = io.StringIO(stdin_text)
|
||
load_runtime_patch = patch.object(main_agent, "load_runtime", return_value=runtime)
|
||
if load_runtime_side_effect is not None:
|
||
load_runtime_patch = patch.object(main_agent, "load_runtime", side_effect=load_runtime_side_effect)
|
||
with load_runtime_patch, \
|
||
patch.object(main_agent, "load_session_index", return_value=session_index), \
|
||
patch.object(sys, "stdin", stdin), \
|
||
contextlib.redirect_stdout(stdout), \
|
||
contextlib.redirect_stderr(stderr):
|
||
code = await main_agent.amain(argv)
|
||
return code, stdout.getvalue(), stderr.getvalue(), runtime
|
||
|
||
def test_help_parser_contains_once(self):
|
||
stdout = io.StringIO()
|
||
with contextlib.redirect_stdout(stdout), self.assertRaises(SystemExit) as ctx:
|
||
main_agent.parse_args(["--help"])
|
||
self.assertEqual(ctx.exception.code, 0)
|
||
self.assertIn("--once", stdout.getvalue())
|
||
|
||
async def test_jsonl_once_outputs_valid_events_with_turn_id(self):
|
||
code, stdout, stderr, runtime = await self.run_cli(["--jsonl", "--once", "hello"])
|
||
self.assertEqual(code, 0)
|
||
self.assertEqual(stderr, "")
|
||
self.assertEqual(runtime.inputs, ["hello"])
|
||
lines = [json.loads(line) for line in stdout.splitlines()]
|
||
self.assertTrue(lines)
|
||
self.assertEqual({event["turn_id"] for event in lines}, {"turn-test"})
|
||
self.assertEqual(lines[0]["type"], "turn")
|
||
self.assertIn("terminal", [event["type"] for event in lines])
|
||
self.assertNotIn("›", stdout)
|
||
|
||
async def test_stdin_non_tty_is_single_prompt(self):
|
||
code, stdout, stderr, runtime = await self.run_cli(["--jsonl"], stdin_text="from stdin\n")
|
||
self.assertEqual(code, 0)
|
||
self.assertEqual(runtime.inputs, ["from stdin"])
|
||
self.assertTrue(stdout.strip())
|
||
|
||
|
||
async def test_stdin_repl_reads_each_line_as_a_turn(self):
|
||
code, stdout, stderr, runtime = await self.run_cli(["--jsonl", "--stdin-repl"], stdin_text="first\nsecond\n")
|
||
self.assertEqual(code, 0)
|
||
self.assertEqual(stderr, "")
|
||
self.assertEqual(runtime.inputs, ["first", "second"])
|
||
events = [json.loads(line) for line in stdout.splitlines()]
|
||
self.assertEqual([event["type"] for event in events if event["type"] == "done"], ["done", "done"])
|
||
|
||
async def test_session_error_is_clear_before_runtime_load(self):
|
||
runtime = FakeRuntime()
|
||
failing_index = FakeSessionIndex(fail_session=True)
|
||
code, stdout, stderr, runtime = await self.run_cli(
|
||
["--session", "missing", "--once", "hello"],
|
||
runtime=runtime,
|
||
session_index=failing_index,
|
||
load_runtime_side_effect=AssertionError("runtime should not load"),
|
||
)
|
||
self.assertEqual(code, 2)
|
||
self.assertEqual(stdout, "")
|
||
self.assertIn("unknown session id", stderr)
|
||
self.assertEqual(runtime.inputs, [])
|
||
|
||
|
||
async def test_new_session_is_created_before_runtime_load(self):
|
||
runtime = FakeRuntime()
|
||
session_index = FakeSessionIndex()
|
||
code, stdout, stderr, runtime = await self.run_cli(
|
||
["--new-session", "--jsonl", "--once", "hello"],
|
||
runtime=runtime,
|
||
session_index=session_index,
|
||
)
|
||
self.assertEqual(code, 0)
|
||
self.assertEqual(stderr, "")
|
||
self.assertEqual(session_index.created, 1)
|
||
self.assertEqual(session_index.activated, "new-1")
|
||
self.assertEqual(runtime.inputs, ["hello"])
|
||
|
||
async def test_new_session_and_session_are_mutually_exclusive(self):
|
||
runtime = FakeRuntime()
|
||
code, stdout, stderr, runtime = await self.run_cli(
|
||
["--new-session", "--session", "s1", "--once", "hello"],
|
||
runtime=runtime,
|
||
load_runtime_side_effect=AssertionError("runtime should not load"),
|
||
)
|
||
self.assertEqual(code, 2)
|
||
self.assertEqual(stdout, "")
|
||
self.assertIn("cannot be used together", stderr)
|
||
self.assertEqual(runtime.inputs, [])
|
||
|
||
|
||
async def test_compact_command_outputs_jsonl_without_llm_turn(self):
|
||
code, stdout, stderr, runtime = await self.run_cli(["--jsonl", "--once", "/compact"])
|
||
self.assertEqual(code, 0)
|
||
self.assertEqual(stderr, "")
|
||
self.assertEqual(runtime.inputs, [])
|
||
events = [json.loads(line) for line in stdout.splitlines()]
|
||
self.assertEqual([event["type"] for event in events], ["turn", "status", "token", "done", "session"])
|
||
turn_ids = {event["turn_id"] for event in events}
|
||
self.assertEqual(len(turn_ids), 1)
|
||
self.assertTrue(next(iter(turn_ids)).startswith("turn-"))
|
||
self.assertIn("Compacted context", events[2]["delta"])
|
||
|
||
async def test_plain_mode_prints_progress_terminal_and_session(self):
|
||
code, stdout, stderr, _ = await self.run_cli(["--once", "hello"])
|
||
self.assertEqual(code, 0)
|
||
self.assertIn("tool: fake_tool running: working", stdout)
|
||
self.assertIn("terminal[stdout]: pipeline output", stdout)
|
||
self.assertIn("session: s1 narrative_planned", stdout)
|