SHA256
130 lines
4.8 KiB
Python
130 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""End-to-end test: fake ollama -> real daemon -> real client agent loop.
|
|
|
|
Exercises the whole chain the same way a user would: the client connects,
|
|
logs in, gets a task, confirms a shell command, and the agent actually
|
|
creates a file in the working directory.
|
|
|
|
Run: python3 test_e2e.py
|
|
"""
|
|
import json
|
|
import os
|
|
import shutil
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
|
|
|
|
def free_ports(count: int) -> list[int]:
|
|
"""Find unused ports by binding ephemeral sockets (avoids TIME_WAIT)."""
|
|
ports, sockets = [], []
|
|
while len(ports) < count:
|
|
s = socket.socket()
|
|
s.bind(("127.0.0.1", 0))
|
|
sockets.append(s)
|
|
ports.append(s.getsockname()[1])
|
|
for s in sockets:
|
|
s.close()
|
|
return ports
|
|
|
|
|
|
class E2E(unittest.TestCase):
|
|
def setUp(self):
|
|
self.ollama_port, self.daemon_port = free_ports(2)
|
|
self.tmp = Path(tempfile.mkdtemp(prefix="sj-e2e-"))
|
|
self.server_dir = self.tmp / "server"
|
|
self.work = self.tmp / "work"
|
|
self.server_dir.mkdir()
|
|
self.work.mkdir()
|
|
(self.server_dir / "config.toml").write_text(
|
|
f'[server]\nhost = "127.0.0.1"\nport = {self.daemon_port}\n'
|
|
f'\n[ollama]\nbase_url = "http://127.0.0.1:{self.ollama_port}"\n'
|
|
'model = "fake-model"\nnum_ctx = 16384\n'
|
|
'\n[database]\npath = "sjcode.db"\n'
|
|
"\n[limits]\nuser = 120\nadmin = -1\n"
|
|
)
|
|
self.prompt_log = self.tmp / "prompts.log"
|
|
|
|
self.fake = subprocess.Popen(
|
|
[sys.executable, str(ROOT / "fake_ollama.py"),
|
|
str(self.ollama_port), str(self.prompt_log)],
|
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
env = dict(os.environ, SJCODE_CONFIG=str(self.server_dir / "config.toml"))
|
|
self.daemon = subprocess.Popen(
|
|
[sys.executable, str(ROOT / "server" / "server.py")],
|
|
cwd=self.server_dir, env=env, stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
|
|
self._wait_ready()
|
|
# create the test account through the daemon admin CLI
|
|
self.daemon.stdin.write("create tester admin\nsecret\nsecret\n")
|
|
self.daemon.stdin.flush()
|
|
time.sleep(0.5) # let the admin CLI process the lines
|
|
|
|
def _wait_ready(self, timeout: float = 10.0) -> None:
|
|
deadline = time.time() + timeout
|
|
url = f"http://127.0.0.1:{self.daemon_port}/status"
|
|
while time.time() < deadline:
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=2) as resp:
|
|
if resp.status == 200:
|
|
return
|
|
except OSError:
|
|
time.sleep(0.1)
|
|
self.fail("daemon did not come up in time")
|
|
|
|
def tearDown(self):
|
|
for proc in (self.daemon, self.fake):
|
|
if proc.poll() is None:
|
|
proc.terminate()
|
|
for proc in (self.daemon, self.fake):
|
|
try:
|
|
proc.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
shutil.rmtree(self.tmp, ignore_errors=True)
|
|
|
|
def _client_run(self, inputs: list[str]) -> subprocess.CompletedProcess:
|
|
script = "\n".join(inputs) + "\n"
|
|
# keep the client's saved-server config inside the test sandbox
|
|
env = dict(os.environ, XDG_CONFIG_HOME=str(self.tmp / "cfg"))
|
|
return subprocess.run(
|
|
[sys.executable, str(ROOT / "client.py"),
|
|
f"127.0.0.1:{self.daemon_port}", "--model", "fake-model"],
|
|
cwd=self.work, env=env, input=script, capture_output=True,
|
|
text=True, timeout=120)
|
|
|
|
def test_full_agent_run_with_confirmation(self):
|
|
proc = self._client_run([
|
|
"tester", # username
|
|
"secret", # password
|
|
"create hello.txt with the word hi in it",
|
|
"y", # confirm the run tool call
|
|
"/quit",
|
|
])
|
|
self.assertEqual(proc.returncode, 0, proc.stdout + proc.stderr)
|
|
self.assertIn("logged in as tester", proc.stdout)
|
|
self.assertIn("Done.", proc.stdout)
|
|
|
|
hello = self.work / "hello.txt"
|
|
self.assertTrue(hello.exists(), "agent did not create hello.txt")
|
|
self.assertEqual(hello.read_text().strip(), "hi")
|
|
|
|
# the model saw the REAL command output, not invented text
|
|
prompts = [json.loads(line)
|
|
for line in self.prompt_log.read_text().splitlines()]
|
|
self.assertEqual(len(prompts), 4) # run, read, done, done-again
|
|
self.assertIn("printf", prompts[1]) # run result fed back
|
|
self.assertIn("Tool results", prompts[1])
|
|
self.assertIn("1\thi", prompts[2]) # read result fed back
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|