test(smolmachines): cover sidecar vm launch paths
This commit is contained in:
@@ -1,11 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import socket
|
||||
import subprocess
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from bot_bottle.backend.smolmachines.port_forward import (
|
||||
ForwardSpec,
|
||||
ForwarderHandle,
|
||||
_CompletedProcess,
|
||||
_Forwarder,
|
||||
_pipe,
|
||||
main,
|
||||
start_forwarder,
|
||||
stop_forwarder,
|
||||
)
|
||||
@@ -30,6 +38,13 @@ def _echo_server() -> tuple[socket.socket, int]:
|
||||
|
||||
|
||||
class TestForwarderValidation(unittest.TestCase):
|
||||
def test_empty_specs_returns_noop_handle(self):
|
||||
handle = start_forwarder(())
|
||||
|
||||
self.assertEqual((), handle.forwards)
|
||||
self.assertEqual(0, handle.process.poll())
|
||||
stop_forwarder(handle)
|
||||
|
||||
def test_rejects_generic_localhost_listener(self):
|
||||
with self.assertRaises(SystemExit):
|
||||
start_forwarder((
|
||||
@@ -66,20 +81,123 @@ class TestForwarderValidation(unittest.TestCase):
|
||||
),
|
||||
))
|
||||
|
||||
def test_rejects_invalid_ports(self):
|
||||
with self.assertRaises(SystemExit):
|
||||
start_forwarder((
|
||||
ForwardSpec("bad-listen", "127.0.0.2", 70000, "127.0.0.1", 9099),
|
||||
))
|
||||
with self.assertRaises(SystemExit):
|
||||
start_forwarder((
|
||||
ForwardSpec("bad-target", "127.0.0.2", 0, "127.0.0.1", 0),
|
||||
))
|
||||
|
||||
def test_completed_process_methods_are_noops(self):
|
||||
proc = _CompletedProcess()
|
||||
|
||||
proc.terminate()
|
||||
proc.kill()
|
||||
|
||||
self.assertEqual(0, proc.poll())
|
||||
self.assertEqual(0, proc.wait(timeout=1))
|
||||
|
||||
|
||||
class _FakeStdout:
|
||||
def __init__(self, line: str):
|
||||
self.line = line
|
||||
self.closed = False
|
||||
|
||||
def readline(self) -> str:
|
||||
return self.line
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class TestForwarderProcess(unittest.TestCase):
|
||||
def test_start_parses_ready_payload(self):
|
||||
payload = {
|
||||
"forwards": [{
|
||||
"label": "egress",
|
||||
"listen_host": "127.0.0.2",
|
||||
"listen_port": 61000,
|
||||
"target_host": "127.0.0.1",
|
||||
"target_port": 55000,
|
||||
}],
|
||||
}
|
||||
proc = MagicMock()
|
||||
proc.stdout = _FakeStdout(json.dumps(payload) + "\n")
|
||||
|
||||
with patch(
|
||||
"bot_bottle.backend.smolmachines.port_forward.subprocess.Popen",
|
||||
return_value=proc,
|
||||
) as popen:
|
||||
handle = start_forwarder((
|
||||
ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 55000),
|
||||
))
|
||||
|
||||
self.assertIs(proc, handle.process)
|
||||
self.assertEqual(61000, handle.forwards[0].listen_port)
|
||||
self.assertTrue(proc.stdout.closed)
|
||||
argv = popen.call_args.args[0]
|
||||
self.assertIn("--spec-json", argv)
|
||||
|
||||
def test_start_dies_when_helper_exits_before_ready(self):
|
||||
proc = MagicMock()
|
||||
proc.stdout = _FakeStdout("")
|
||||
|
||||
with patch(
|
||||
"bot_bottle.backend.smolmachines.port_forward.subprocess.Popen",
|
||||
return_value=proc,
|
||||
):
|
||||
with self.assertRaises(SystemExit):
|
||||
start_forwarder((
|
||||
ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 55000),
|
||||
))
|
||||
|
||||
proc.terminate.assert_called_once()
|
||||
proc.wait.assert_called_once_with(timeout=5)
|
||||
|
||||
def test_start_dies_on_invalid_ready_payload(self):
|
||||
proc = MagicMock()
|
||||
proc.stdout = _FakeStdout('{"wrong": []}\n')
|
||||
|
||||
with patch(
|
||||
"bot_bottle.backend.smolmachines.port_forward.subprocess.Popen",
|
||||
return_value=proc,
|
||||
):
|
||||
with self.assertRaises(SystemExit):
|
||||
start_forwarder((
|
||||
ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 55000),
|
||||
))
|
||||
|
||||
proc.terminate.assert_called_once()
|
||||
proc.wait.assert_called_once_with(timeout=5)
|
||||
|
||||
def test_stop_kills_process_after_timeout(self):
|
||||
proc = MagicMock()
|
||||
proc.poll.return_value = None
|
||||
proc.wait.side_effect = [subprocess.TimeoutExpired("cmd", 5), 0]
|
||||
|
||||
stop_forwarder(ForwarderHandle(process=proc, forwards=()))
|
||||
|
||||
proc.terminate.assert_called_once()
|
||||
proc.kill.assert_called_once()
|
||||
self.assertEqual(2, proc.wait.call_count)
|
||||
|
||||
|
||||
class TestForwarding(unittest.TestCase):
|
||||
def test_forwards_bytes_both_directions(self):
|
||||
try:
|
||||
probe = socket.create_server(("127.0.0.16", 0))
|
||||
probe = socket.create_server(("127.0.0.2", 0))
|
||||
except OSError:
|
||||
self.skipTest("127.0.0.16 is not bound on this host")
|
||||
self.skipTest("127.0.0.2 is not bindable on this host")
|
||||
else:
|
||||
probe.close()
|
||||
_listener, target_port = _echo_server()
|
||||
handle = start_forwarder((
|
||||
ForwardSpec(
|
||||
label="egress",
|
||||
listen_host="127.0.0.16",
|
||||
listen_host="127.0.0.2",
|
||||
listen_port=0,
|
||||
target_host="127.0.0.1",
|
||||
target_port=target_port,
|
||||
@@ -94,5 +212,114 @@ class TestForwarding(unittest.TestCase):
|
||||
stop_forwarder(handle)
|
||||
|
||||
|
||||
class TestForwarderInProcess(unittest.TestCase):
|
||||
def test_forwarder_serves_and_stops_in_process(self):
|
||||
_listener, target_port = _echo_server()
|
||||
forwarder = _Forwarder((
|
||||
ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", target_port),
|
||||
))
|
||||
forwarder.start()
|
||||
thread = threading.Thread(target=forwarder.serve, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
bound = forwarder.bound[0]
|
||||
with socket.create_connection((bound.listen_host, bound.listen_port)) as sock:
|
||||
sock.sendall(b"ping")
|
||||
self.assertEqual(b"ping", sock.recv(4))
|
||||
finally:
|
||||
forwarder.stop()
|
||||
thread.join(timeout=2)
|
||||
|
||||
self.assertFalse(thread.is_alive())
|
||||
|
||||
def test_forwarder_logs_target_connect_failure(self):
|
||||
forwarder = _Forwarder((
|
||||
ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 1),
|
||||
))
|
||||
forwarder.start()
|
||||
thread = threading.Thread(target=forwarder.serve, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
bound = forwarder.bound[0]
|
||||
with patch("bot_bottle.backend.smolmachines.port_forward.warn") as warn:
|
||||
with socket.create_connection((bound.listen_host, bound.listen_port)):
|
||||
pass
|
||||
for _ in range(100):
|
||||
if warn.called:
|
||||
break
|
||||
threading.Event().wait(0.01)
|
||||
self.assertTrue(warn.called)
|
||||
finally:
|
||||
forwarder.stop()
|
||||
thread.join(timeout=2)
|
||||
|
||||
def test_pipe_stops_when_recv_fails(self):
|
||||
src, dst = socket.socketpair()
|
||||
src.close()
|
||||
try:
|
||||
_pipe(src, dst)
|
||||
finally:
|
||||
dst.close()
|
||||
|
||||
def test_pipe_stops_when_send_fails(self):
|
||||
src, src_peer = socket.socketpair()
|
||||
_dst_peer, dst = socket.socketpair()
|
||||
dst.close()
|
||||
try:
|
||||
src_peer.sendall(b"hello")
|
||||
src_peer.shutdown(socket.SHUT_WR)
|
||||
_pipe(src, dst)
|
||||
finally:
|
||||
src.close()
|
||||
src_peer.close()
|
||||
_dst_peer.close()
|
||||
|
||||
|
||||
class TestMain(unittest.TestCase):
|
||||
def test_main_starts_serves_and_stops_forwarder(self):
|
||||
calls: list[str] = []
|
||||
handlers = []
|
||||
|
||||
class FakeForwarder:
|
||||
def __init__(self, specs): # type: ignore[no-untyped-def]
|
||||
self.specs = specs
|
||||
self.bound = specs
|
||||
|
||||
def start(self) -> None:
|
||||
calls.append("start")
|
||||
|
||||
def serve(self) -> None:
|
||||
calls.append("serve")
|
||||
|
||||
def stop(self) -> None:
|
||||
calls.append("stop")
|
||||
|
||||
def record_signal(_signum, handler): # type: ignore[no-untyped-def]
|
||||
handlers.append(handler)
|
||||
|
||||
spec = ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 9099)
|
||||
with patch(
|
||||
"bot_bottle.backend.smolmachines.port_forward._Forwarder",
|
||||
FakeForwarder,
|
||||
), patch(
|
||||
"bot_bottle.backend.smolmachines.port_forward.signal.signal",
|
||||
side_effect=record_signal,
|
||||
), patch("builtins.print") as print_mock:
|
||||
self.assertEqual(
|
||||
0,
|
||||
main(("--spec-json", json.dumps([{
|
||||
"label": spec.label,
|
||||
"listen_host": spec.listen_host,
|
||||
"listen_port": spec.listen_port,
|
||||
"target_host": spec.target_host,
|
||||
"target_port": spec.target_port,
|
||||
}]))),
|
||||
)
|
||||
handlers[0](15, None)
|
||||
|
||||
self.assertEqual(["start", "serve", "stop", "stop"], calls)
|
||||
self.assertIn("forwards", print_mock.call_args.args[0])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user