SHA256
278 lines
11 KiB
Python
278 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Tests for the sjcode client agent layer.
|
|
|
|
Run: python3 -m unittest test_client.py -v
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
import client as c # noqa: E402
|
|
|
|
|
|
class ClientTestCase(unittest.TestCase):
|
|
"""Each test runs in its own temp dir which becomes the project dir."""
|
|
|
|
def setUp(self) -> None:
|
|
self._old_cwd = c.CWD
|
|
self._old_chdir = os.getcwd()
|
|
self._tmp = tempfile.TemporaryDirectory()
|
|
c.CWD = Path(self._tmp.name).resolve()
|
|
os.chdir(c.CWD) # tools use CWD; keep relative paths consistent
|
|
|
|
def tearDown(self) -> None:
|
|
os.chdir(self._old_chdir)
|
|
c.CWD = self._old_cwd
|
|
self._tmp.cleanup()
|
|
|
|
|
|
class TestParsing(ClientTestCase):
|
|
def test_extracts_run_and_json_tools_in_order(self) -> None:
|
|
raw = (
|
|
"Let me check the repo.\n"
|
|
"```run\n"
|
|
"git status --short\n"
|
|
"```\n"
|
|
"And now the file:\n"
|
|
"```sj:read\n"
|
|
'{"path": "a.py", "limit": 10}\n'
|
|
"```\n"
|
|
)
|
|
tags = c.read_tool_tags(raw)
|
|
self.assertEqual([("run", "git status --short"),
|
|
("read", '{"path": "a.py", "limit": 10}')], tags)
|
|
|
|
def test_language_fence_is_run_tool(self) -> None:
|
|
self.assertEqual(c.read_tool_tags("```run\necho hi\n```"),
|
|
[("run", "echo hi")])
|
|
|
|
def test_language_fence_is_not_a_run_tool(self) -> None:
|
|
self.assertEqual(c.read_tool_tags("```rust\nfn main() {}\n```"), [])
|
|
|
|
def test_text_before_and_after_is_stripped_from_reply(self) -> None:
|
|
raw = "Thinking...\n```run\necho hi\n```\nDone."
|
|
had, reply, results = c.process_reply(raw, auto=True)
|
|
self.assertTrue(had)
|
|
self.assertNotIn("```", reply)
|
|
self.assertNotIn("echo hi", reply)
|
|
|
|
def test_body_with_triple_backticks_inside_json_survives(self) -> None:
|
|
body = '{"path": "x.md", "content": "code: ```py\\nprint(1)\\n```\\n"}'
|
|
raw = f"```sj:write\n{body}\n```\n"
|
|
self.assertEqual(c.read_tool_tags(raw), [("write", body)])
|
|
|
|
def test_unclosed_fence_still_captured(self) -> None:
|
|
self.assertEqual(c.read_tool_tags("```run\necho hi"),
|
|
[("run", "echo hi")])
|
|
|
|
|
|
class TestRunTool(ClientTestCase):
|
|
def test_runs_command_and_returns_output(self) -> None:
|
|
ok, lines = c.handle_tool_call("run", '{"command": "echo hello"}', auto=True)
|
|
self.assertTrue(ok)
|
|
self.assertIn("hello", "\n".join(lines))
|
|
|
|
def test_nonzero_exit_reported(self) -> None:
|
|
ok, lines = c.handle_tool_call(
|
|
"run", '{"command": "python3 -c \x27import sys; sys.exit(3)\x27"}',
|
|
auto=True)
|
|
self.assertFalse(ok)
|
|
self.assertIn("exit 3", "\n".join(lines))
|
|
|
|
def test_cwd_argument_scopes_the_command(self) -> None:
|
|
os.mkdir("sub")
|
|
Path("sub/marker.txt").write_text("x")
|
|
ok, lines = c.handle_tool_call(
|
|
"run", '{"command": "pwd", "cwd": "sub"}', auto=True)
|
|
self.assertTrue(ok)
|
|
self.assertIn("sub", "\n".join(lines))
|
|
|
|
def test_run_outside_project_dir_is_refused(self) -> None:
|
|
ok, lines = c.handle_tool_call(
|
|
"run", json.dumps({"command": "pwd", "cwd": "/tmp"}), auto=True)
|
|
self.assertFalse(ok)
|
|
self.assertIn("outside the project directory", "\n".join(lines))
|
|
|
|
def test_missing_command_is_refused(self) -> None:
|
|
ok, lines = c.handle_tool_call("run", "{}", auto=True)
|
|
self.assertFalse(ok)
|
|
self.assertIn("command", "\n".join(lines))
|
|
|
|
|
|
class TestFileTools(ClientTestCase):
|
|
def test_read_shows_numbered_lines(self) -> None:
|
|
Path("a.txt").write_text("one\ntwo\nthree\n")
|
|
ok, lines = c.handle_tool_call("read", '{"path": "a.txt"}', auto=True)
|
|
self.assertTrue(ok)
|
|
joined = "\n".join(lines)
|
|
self.assertIn("1\tone", joined)
|
|
self.assertIn("3\tthree", joined)
|
|
|
|
def test_read_supports_start_and_limit(self) -> None:
|
|
Path("a.txt").write_text("\n".join(str(i) for i in range(1, 51)))
|
|
ok, lines = c.handle_tool_call(
|
|
"read", '{"path": "a.txt", "start": 40, "limit": 5}', auto=True)
|
|
self.assertIn("40\t40", "\n".join(lines))
|
|
|
|
def test_read_missing_file_gives_error_back(self) -> None:
|
|
ok, lines = c.handle_tool_call("read", '{"path": "nope.txt"}', auto=True)
|
|
self.assertFalse(ok)
|
|
self.assertIn("not a file", "\n".join(lines))
|
|
|
|
def test_files_lists_and_skips_junk_dirs(self) -> None:
|
|
os.mkdir("pkg")
|
|
Path("pkg/mod.py").write_text("x = 1\n")
|
|
os.mkdir(".git")
|
|
Path(".git/config").write_text("x\n")
|
|
ok, lines = c.handle_tool_call("files", '{"path": "."}', auto=True)
|
|
joined = "\n".join(lines)
|
|
self.assertIn("pkg/mod.py", joined)
|
|
self.assertNotIn(".git/config", joined)
|
|
|
|
def test_write_creates_file_after_confirmation(self) -> None:
|
|
with mock.patch("client.confirm_tool", return_value="y"):
|
|
ok, lines = c.handle_tool_call(
|
|
"write", '{"path": "new.txt", "content": "hi\\n"}', auto=False)
|
|
self.assertTrue(ok)
|
|
self.assertEqual(Path("new.txt").read_text(), "hi\n")
|
|
|
|
def test_write_refused_outside_project(self) -> None:
|
|
ok, lines = c.handle_tool_call(
|
|
"write",
|
|
json.dumps({"path": "../../evil.txt", "content": "x"}),
|
|
auto=True)
|
|
self.assertFalse(ok)
|
|
self.assertIn("outside the project directory", "\n".join(lines))
|
|
|
|
def test_edit_replaces_exact_text(self) -> None:
|
|
Path("app.py").write_text("def main():\n print('old')\n")
|
|
with mock.patch("client.confirm_tool", return_value="y"):
|
|
ok, lines = c.handle_tool_call(
|
|
"edit",
|
|
json.dumps({"path": "app.py", "old": "print('old')",
|
|
"new": "print('new')"}),
|
|
auto=False)
|
|
self.assertTrue(ok)
|
|
self.assertIn("print('new')", Path("app.py").read_text())
|
|
|
|
def test_edit_requires_unique_match_without_all(self) -> None:
|
|
Path("app.py").write_text("x = 1\nx = 1\n")
|
|
ok, lines = c.handle_tool_call(
|
|
"edit", json.dumps({"path": "app.py", "old": "x = 1", "new": "y"}),
|
|
auto=True)
|
|
self.assertFalse(ok)
|
|
self.assertIn("matches 2 times", "\n".join(lines))
|
|
|
|
def test_edit_with_all_flag_replaces_every_match(self) -> None:
|
|
Path("app.py").write_text("x = 1\nx = 1\n")
|
|
with mock.patch("client.confirm_tool", return_value="y"):
|
|
ok, lines = c.handle_tool_call(
|
|
"edit",
|
|
json.dumps({"path": "app.py", "old": "x = 1", "new": "y",
|
|
"all": True}),
|
|
auto=False)
|
|
self.assertTrue(ok)
|
|
self.assertEqual(Path("app.py").read_text(), "y\ny\n")
|
|
|
|
def test_edit_not_found_gives_actionable_error(self) -> None:
|
|
Path("app.py").write_text("hello\n")
|
|
ok, lines = c.handle_tool_call(
|
|
"edit", json.dumps({"path": "app.py", "old": "zzz", "new": "y"}),
|
|
auto=True)
|
|
self.assertFalse(ok)
|
|
self.assertIn("search text not found", "\n".join(lines))
|
|
|
|
def test_overwrite_of_big_existing_file_refused(self) -> None:
|
|
Path("big.txt").write_text("x" * 3000)
|
|
ok, lines = c.handle_tool_call(
|
|
"write", json.dumps({"path": "big.txt", "content": "y"}), auto=True)
|
|
self.assertFalse(ok)
|
|
self.assertIn("use 'edit' instead", "\n".join(lines))
|
|
|
|
|
|
class TestConfirmation(ClientTestCase):
|
|
def test_declined_call_returns_cancelled_to_model(self) -> None:
|
|
Path("f.txt").write_text("a\n")
|
|
with mock.patch("client.confirm_tool", return_value="n"):
|
|
ok, lines = c.handle_tool_call(
|
|
"edit", json.dumps({"path": "f.txt", "old": "a", "new": "b"}),
|
|
auto=False)
|
|
self.assertFalse(ok)
|
|
self.assertIn("cancelled by user", "\n".join(lines))
|
|
|
|
def test_user_instruction_is_passed_to_model(self) -> None:
|
|
Path("f.txt").write_text("a\n")
|
|
with mock.patch("client.confirm_tool",
|
|
return_value="use python instead"):
|
|
ok, lines = c.handle_tool_call("run", '{"command": "ls -la"}',
|
|
auto=False)
|
|
self.assertFalse(ok)
|
|
self.assertIn("user says: use python instead", "\n".join(lines))
|
|
|
|
def test_abort_raises_toolabort(self) -> None:
|
|
with mock.patch("client.confirm_tool", return_value=c.ABORT):
|
|
with self.assertRaises(c.ToolAbort):
|
|
c.handle_tool_call("run", '{"command": "ls"}', auto=False)
|
|
|
|
def test_auto_mode_skips_confirmation(self) -> None:
|
|
with mock.patch("client.confirm_tool",
|
|
side_effect=AssertionError("must not be called")):
|
|
ok, lines = c.handle_tool_call("run", '{"command": "echo ok"}',
|
|
auto=True)
|
|
self.assertTrue(ok)
|
|
|
|
|
|
class TestAgentTurn(ClientTestCase):
|
|
def _chat(self, replies):
|
|
"""Patch api_chat to serve canned replies in order."""
|
|
calls = {"n": 0}
|
|
|
|
def fake(entry, token, message, model, on_chunk=None):
|
|
i = calls["n"]
|
|
calls["n"] += 1
|
|
return 200, {"response": replies[min(i, len(replies) - 1)],
|
|
"remaining_minutes": 99}
|
|
|
|
return mock.patch("client.api_chat", side_effect=fake)
|
|
|
|
def test_tool_loop_feeds_results_and_finishes_on_done(self) -> None:
|
|
Path("greet.txt").write_text("hi\n")
|
|
replies = [
|
|
'Reading the file.\n```sj:read\n{"path": "greet.txt"}\n```',
|
|
'```sj:write\n{"path": "out.txt", "content": "done"}\n```',
|
|
"Everything is in place.\nDone.",
|
|
]
|
|
with self._chat(replies):
|
|
status, final, _ = c.agent_turn({}, "t", "m", [], "make out.txt",
|
|
auto=True)
|
|
self.assertEqual(status, "done")
|
|
self.assertEqual(Path("out.txt").read_text(), "done\n")
|
|
self.assertEqual(final, "Everything is in place.\nDone.")
|
|
|
|
def test_plain_text_reply_ends_turn_without_tool_loop(self) -> None:
|
|
with self._chat(["Just answering."]):
|
|
status, final, _ = c.agent_turn({}, "t", "m", [], "hello", auto=True)
|
|
self.assertEqual(status, "done")
|
|
self.assertEqual(final, "Just answering.")
|
|
|
|
def test_error_result_is_sent_back_to_model(self) -> None:
|
|
replies = [
|
|
'```sj:read\n{"path": "missing.txt"}\n```',
|
|
"Task complete.\nDone.",
|
|
]
|
|
with self._chat(replies):
|
|
status, _, _ = c.agent_turn({}, "t", "m", [], "read it", auto=True)
|
|
self.assertEqual(status, "done")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|