120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.util
|
|
from pathlib import Path
|
|
import socket
|
|
import sys
|
|
import threading
|
|
import unittest
|
|
|
|
|
|
SCRIPTS = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(SCRIPTS))
|
|
|
|
from port_forwarding_benchmark_common import endpoint_measurement, summarize
|
|
|
|
|
|
def load_script(name: str):
|
|
path = SCRIPTS / name
|
|
spec = importlib.util.spec_from_file_location(path.stem, path)
|
|
assert spec is not None and spec.loader is not None
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
termind_benchmark = load_script("benchmark-port-forwarding.py")
|
|
external_benchmark = load_script("benchmark-external-port-forwarding.py")
|
|
|
|
|
|
class TermindArgumentsTests(unittest.TestCase):
|
|
def test_default_and_unique_connection_counts(self) -> None:
|
|
arguments = termind_benchmark.make_parser().parse_args(
|
|
["--metadata", "fixture.json", "--output", "result.json"]
|
|
)
|
|
self.assertEqual(arguments.termind_connections, [1])
|
|
self.assertEqual(termind_benchmark.unique_counts([1, 64, 1, 4]), [1, 64, 4])
|
|
|
|
def test_connection_count_accepts_boundaries(self) -> None:
|
|
self.assertEqual(termind_benchmark.connection_count("1"), 1)
|
|
self.assertEqual(termind_benchmark.connection_count("64"), 64)
|
|
|
|
def test_connection_count_rejects_invalid_values(self) -> None:
|
|
for value in ("0", "65", "not-a-number"):
|
|
with self.subTest(value=value):
|
|
with self.assertRaises(argparse.ArgumentTypeError):
|
|
termind_benchmark.connection_count(value)
|
|
|
|
|
|
class ExternalArgumentsTests(unittest.TestCase):
|
|
def test_system_and_manual_modes_are_distinct(self) -> None:
|
|
parser = external_benchmark.make_parser()
|
|
system = parser.parse_args(
|
|
["--metadata", "fixture.json", "--output", "result.json", "system-ssh"]
|
|
)
|
|
manual = parser.parse_args(
|
|
[
|
|
"--metadata",
|
|
"fixture.json",
|
|
"--output",
|
|
"result.json",
|
|
"external",
|
|
"--name",
|
|
"Example",
|
|
"--kind",
|
|
"dynamic",
|
|
"--port",
|
|
"1080",
|
|
]
|
|
)
|
|
self.assertEqual(system.mode, "system-ssh")
|
|
self.assertEqual((manual.mode, manual.kind, manual.port), ("external", "dynamic", 1080))
|
|
|
|
|
|
class WorkloadTests(unittest.TestCase):
|
|
def test_manual_endpoint_reports_network_metrics_only(self) -> None:
|
|
listener = socket.socket()
|
|
listener.bind(("127.0.0.1", 0))
|
|
listener.listen()
|
|
stopped = threading.Event()
|
|
|
|
def serve() -> None:
|
|
while not stopped.is_set():
|
|
try:
|
|
connection, _ = listener.accept()
|
|
except OSError:
|
|
return
|
|
threading.Thread(target=echo, args=(connection,), daemon=True).start()
|
|
|
|
def echo(connection: socket.socket) -> None:
|
|
with connection:
|
|
while data := connection.recv(65_536):
|
|
connection.sendall(data)
|
|
|
|
thread = threading.Thread(target=serve, daemon=True)
|
|
thread.start()
|
|
try:
|
|
record = endpoint_measurement(
|
|
"external:Example",
|
|
listener.getsockname()[1],
|
|
"local",
|
|
streams=2,
|
|
target_port=0,
|
|
rounds=1,
|
|
)
|
|
finally:
|
|
stopped.set()
|
|
listener.close()
|
|
self.assertEqual(record["ssh_connections"], 1)
|
|
self.assertEqual(record["jumps"], 0)
|
|
self.assertNotIn("cpu_s", record)
|
|
self.assertNotIn("forwarder_peak_rss_mib", record)
|
|
summary = summarize([record])
|
|
self.assertEqual(summary[0]["samples"], 1)
|
|
self.assertNotIn("median_cpu_s", summary[0])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|