93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
from beaver_gateway.backends.processes import Processes, _command, subtree
|
|
|
|
SLEEP = "import time; time.sleep(60)"
|
|
NEST = (
|
|
"import subprocess, sys, time; "
|
|
f"subprocess.Popen([sys.executable, '-c', {SLEEP!r}]); time.sleep(60)"
|
|
)
|
|
|
|
|
|
def spawn(code: str) -> subprocess.Popen[bytes]:
|
|
return subprocess.Popen([sys.executable, "-c", code])
|
|
|
|
|
|
def rows_under(processes: Processes, pid: int, count: int) -> list[dict]:
|
|
deadline = time.monotonic() + 5
|
|
while time.monotonic() < deadline:
|
|
rows = subtree(processes.tree_parents_first(), pid)
|
|
if len(rows) >= count:
|
|
return rows
|
|
time.sleep(0.05)
|
|
msg = f"{pid} never grew to {count} processes"
|
|
raise AssertionError(msg)
|
|
|
|
|
|
def test_the_tree_shows_what_the_gateway_started_and_what_it_costs() -> None:
|
|
child = spawn(SLEEP)
|
|
try:
|
|
processes = Processes()
|
|
(row,) = rows_under(processes, child.pid, 1)
|
|
assert row["parent"] == os.getpid()
|
|
assert row["depth"] == 0
|
|
assert row["rss"] > 0
|
|
assert row["cpu"] is None
|
|
assert "time.sleep" in row["command"]
|
|
(again,) = rows_under(processes, child.pid, 1)
|
|
assert again["cpu"] is not None
|
|
finally:
|
|
child.kill()
|
|
child.wait()
|
|
|
|
|
|
def test_kill_takes_the_branch_and_nothing_the_gateway_did_not_start() -> None:
|
|
parent = spawn(NEST)
|
|
try:
|
|
processes = Processes()
|
|
top, nested = rows_under(processes, parent.pid, 2)
|
|
assert (top["depth"], nested["depth"]) == (0, 1)
|
|
assert not processes.kill_own_descendant(os.getpid())
|
|
assert not processes.kill_own_descendant(os.getppid())
|
|
assert processes.kill_own_descendant(parent.pid)
|
|
parent.wait(timeout=5)
|
|
alive = {row["pid"] for row in processes.tree_parents_first()}
|
|
assert not alive & {top["pid"], nested["pid"]}
|
|
finally:
|
|
parent.kill()
|
|
parent.wait()
|
|
|
|
|
|
def test_a_subtree_ends_at_the_next_sibling() -> None:
|
|
rows = [
|
|
{"pid": 1, "depth": 0},
|
|
{"pid": 2, "depth": 1},
|
|
{"pid": 3, "depth": 2},
|
|
{"pid": 4, "depth": 1},
|
|
{"pid": 5, "depth": 0},
|
|
]
|
|
assert [(r["pid"], r["depth"]) for r in subtree(rows, 2)] == [(2, 0), (3, 1)]
|
|
assert [r["pid"] for r in subtree(rows, 1)] == [1, 2, 3, 4]
|
|
assert subtree(rows, 9) == []
|
|
assert subtree(rows, None) == []
|
|
|
|
|
|
def test_a_command_line_keeps_prompts_and_configs_to_itself() -> None:
|
|
argv = [
|
|
"/opt/claude/claude",
|
|
"--model",
|
|
"m",
|
|
"--system-prompt",
|
|
"you are a beaver",
|
|
"--mcp-config",
|
|
'{"gateway": {"headers": {"Authorization": "Bearer abc"}}}',
|
|
"x" * 100,
|
|
]
|
|
assert _command(argv) == (
|
|
"claude --model m --system-prompt … --mcp-config … " + "x" * 80 + "…"
|
|
)
|
|
assert _command([]) == ""
|