import asyncio
import json
import os
import socket
import sys
import tempfile
import types
import unittest
import urllib.error
import urllib.request
from unittest import mock
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "files"))
import claude_agents as ca # noqa: E402
def free_port():
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def http_get(port, path):
with urllib.request.urlopen(f"http://127.0.0.1:{port}{path}", timeout=3) as r:
return r.status, r.read().decode("utf-8")
class PortAlive(unittest.TestCase):
def test_port_alive(self):
s = socket.socket()
s.bind(("127.0.0.1", 0))
s.listen()
port = s.getsockname()[1]
self.assertTrue(ca.port_alive(port))
s.close()
self.assertFalse(ca.port_alive(port))
class Serve(unittest.TestCase):
def setUp(self):
self.port = free_port()
patches = [
mock.patch.object(ca, "PORT", self.port),
mock.patch.object(ca, "CONFIG_PATH", "/nonexistent/claude-agents/config.json"),
mock.patch.object(ca, "LOG_DIR", self._tmpdir()),
mock.patch.object(ca, "load_sessions", lambda: []),
mock.patch.object(ca, "fetch_usage", lambda: {}),
mock.patch.object(ca, "fetch_balance", lambda: ({}, "test")),
]
for p in patches:
p.start()
self.addCleanup(p.stop)
def _tmpdir(self):
d = tempfile.TemporaryDirectory()
self.addCleanup(d.cleanup)
return d.name
async def _with_server(self, check):
task = asyncio.create_task(ca.serve())
loop = asyncio.get_running_loop()
try:
for _ in range(60):
if ca.port_alive(self.port):
break
await asyncio.sleep(0.05)
return await check(loop)
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
def test_serve_answers(self):
async def check(loop):
status, body = await loop.run_in_executor(None, http_get, self.port, "/")
self.assertEqual(status, 200)
self.assertIn('
', body)
status, raw = await loop.run_in_executor(None, http_get, self.port, "/raw")
self.assertIn("tree", json.loads(raw))
asyncio.run(self._with_server(check))
self.assertFalse(ca.port_alive(self.port))
def test_serve_survives_usage_failure(self):
def no_token():
raise FileNotFoundError("no credentials")
async def check(loop):
body = ""
for _ in range(60):
_, body = await loop.run_in_executor(None, http_get, self.port, "/limits")
if "Не найден токен Claude Code" in body:
break
await asyncio.sleep(0.05)
self.assertIn("Не найден токен Claude Code", body)
status, _ = await loop.run_in_executor(None, http_get, self.port, "/")
self.assertEqual(status, 200)
with mock.patch.object(ca, "fetch_usage", no_token):
asyncio.run(self._with_server(check))
def test_host_header_checked(self):
def get(host, path="/", **extra):
req = urllib.request.Request(f"http://127.0.0.1:{self.port}{path}", headers={"Host": host, **extra})
try:
with urllib.request.urlopen(req, timeout=3) as r:
return r.status
except urllib.error.HTTPError as e:
e.close()
return e.code
async def check(loop):
self.assertEqual(await loop.run_in_executor(None, get, "evil.example:%d" % self.port), 403)
self.assertEqual(await loop.run_in_executor(None, get, "127.0.0.1"), 403)
self.assertEqual(await loop.run_in_executor(None, get, "127.0.0.1:%d" % self.port), 200)
self.assertEqual(await loop.run_in_executor(None, get, "LocalHost:%d" % self.port), 200)
self.assertEqual(await loop.run_in_executor(None, get, "evil.example:%d" % self.port, "/raw"), 403)
asyncio.run(self._with_server(check))
def test_cross_site_actions_forbidden(self):
def get(path, site):
req = urllib.request.Request(f"http://127.0.0.1:{self.port}{path}", headers={"Sec-Fetch-Site": site})
try:
with urllib.request.urlopen(req, timeout=3) as r:
return r.status
except urllib.error.HTTPError as e:
e.close()
return e.code
async def check(loop):
for path in ("/activate?id=x", "/open-wt?id=x&path=/tmp"):
self.assertEqual(await loop.run_in_executor(None, get, path, "cross-site"), 403)
with mock.patch.object(ca, "activate", mock.AsyncMock()) as act, \
mock.patch.object(ca, "open_worktree", mock.AsyncMock()) as wt:
asyncio.run(self._with_server(check))
act.assert_not_called()
wt.assert_not_called()
def test_serve_busy_port_raises(self):
s = socket.socket()
s.bind(("127.0.0.1", self.port))
s.listen()
self.addCleanup(s.close)
with self.assertRaises(OSError):
asyncio.run(ca.serve())
class RegisterToolbelt(unittest.TestCase):
def setUp(self):
self.registered = []
self.kicks = []
async def register(connection, **kwargs):
self.registered.append(kwargs)
fake = types.SimpleNamespace(tool=types.SimpleNamespace(async_register_web_view_tool=register))
patches = [
mock.patch.object(ca, "iterm2", fake),
mock.patch.object(ca, "PORT", 8766),
mock.patch.object(ca, "CONFIG_PATH", "/nonexistent/claude-agents/config.json"),
mock.patch.object(ca, "kickstart_service", lambda: self.kicks.append(1)),
]
for p in patches:
p.start()
self.addCleanup(p.stop)
def assert_registered(self):
self.assertEqual(self.registered, [{
"display_name": "Claude", "identifier": "com.andrey.claude-agents",
"reveal_if_already_registered": False, "url": "http://127.0.0.1:8766/"}])
def test_register_toolbelt_kickstarts_when_port_dead(self):
async def no_sleep(sec):
pass
with mock.patch.object(ca, "port_alive", lambda *a, **k: False), \
mock.patch.object(ca.asyncio, "sleep", no_sleep):
asyncio.run(ca.register_toolbelt(object()))
self.assertEqual(self.kicks, [1])
self.assert_registered()
def test_register_toolbelt_waits_for_service(self):
alive = iter([False, False, False, True]) # проверка в начале + ожидание после kickstart
sleeps = []
async def fake_sleep(sec):
sleeps.append(sec)
with mock.patch.object(ca, "port_alive", lambda *a, **k: next(alive)), \
mock.patch.object(ca.asyncio, "sleep", fake_sleep):
asyncio.run(ca.register_toolbelt(object()))
self.assertEqual(sleeps, [0.25, 0.25])
self.assert_registered()
def test_register_toolbelt_skips_kickstart_when_alive(self):
with mock.patch.object(ca, "port_alive", lambda *a, **k: True):
asyncio.run(ca.register_toolbelt(object()))
self.assertEqual(self.kicks, [])
self.assert_registered()
class Kickstart(unittest.TestCase):
def test_kickstart_command(self):
with mock.patch.object(ca.subprocess, "run") as run:
ca.kickstart_service()
self.assertEqual(run.call_args.args[0],
["launchctl", "kickstart", f"gui/{os.getuid()}/com.andrey.claude-agents"])
with mock.patch.object(ca.subprocess, "run", side_effect=OSError("no launchctl")):
ca.kickstart_service() # не бросает
class Run(unittest.TestCase):
def test_run_headless_calls_serve(self):
called = []
async def fake_serve():
called.append(1)
with mock.patch.object(ca, "serve", fake_serve):
ca.run(["--headless"])
self.assertEqual(called, [1])
def test_run_without_iterm_exits(self):
with mock.patch.object(ca, "iterm2", None):
with self.assertRaises(SystemExit):
ca.run([])
if __name__ == "__main__":
unittest.main()