import http.server import io import json import socket import threading import unittest from latprobe.cli import ( EXIT_CONNECT, EXIT_DNS, EXIT_HTTP, EXIT_OK, EXIT_TIMEOUT, EXIT_USAGE, run, ) class _OKHandler(http.server.BaseHTTPRequestHandler): def do_GET(self): self.send_response(200) self.end_headers() self.wfile.write(b"hello latprobe") def log_message(self, *args): pass class _NotFoundHandler(http.server.BaseHTTPRequestHandler): def do_GET(self): self.send_response(404) self.end_headers() self.wfile.write(b"not found") def log_message(self, *args): pass def _start_server(handler_class): server = http.server.HTTPServer(("127.0.0.1", 0), handler_class) t = threading.Thread(target=server.serve_forever) t.daemon = True t.start() return server, server.server_address[1] def _free_port() -> int: with socket.socket() as s: s.bind(("127.0.0.1", 0)) return s.getsockname()[1] def _black_hole_port() -> int: srv = socket.socket() srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) srv.bind(("127.0.0.1", 0)) srv.listen(10) port = srv.getsockname()[1] conns: list = [] def _serve(): while True: try: conn, _ = srv.accept() conns.append(conn) except OSError: break threading.Thread(target=_serve, daemon=True).start() return port def _invoke(args: list[str]) -> tuple[int, str, str]: out, err = io.StringIO(), io.StringIO() code = run(args, out, err) return code, out.getvalue(), err.getvalue() class TestCLIUsageErrors(unittest.TestCase): def test_no_args_returns_usage(self): code, out, err = _invoke([]) self.assertEqual(code, EXIT_USAGE) def test_help_returns_ok(self): code, out, err = _invoke(["-h"]) self.assertEqual(code, EXIT_OK) self.assertIn("latprobe", out) def test_invalid_timeout_returns_usage(self): code, out, err = _invoke(["--timeout", "bad", "http://example.com"]) self.assertEqual(code, EXIT_USAGE) self.assertIn("timeout", err) class TestCLISuccess(unittest.TestCase): @classmethod def setUpClass(cls): cls.ok_server, cls.ok_port = _start_server(_OKHandler) cls.nf_server, cls.nf_port = _start_server(_NotFoundHandler) @classmethod def tearDownClass(cls): cls.ok_server.shutdown() cls.nf_server.shutdown() def _url(self, port=None): return f"http://127.0.0.1:{port or self.ok_port}" def test_single_url_success(self): code, out, err = _invoke([self._url()]) self.assertEqual(code, EXIT_OK) self.assertIn("200", out) self.assertIn("Total", out) self.assertIn("DNS lookup", out) def test_output_has_separator(self): code, out, _ = _invoke([self._url()]) self.assertIn("─", out) def test_count_shows_aggregate(self): code, out, _ = _invoke(["-n", "3", self._url()]) self.assertEqual(code, EXIT_OK) self.assertIn("3 samples", out) self.assertIn("min", out) self.assertIn("avg", out) self.assertIn("max", out) def test_multiple_urls_output_separated_by_blank_line(self): url = self._url() code, out, _ = _invoke([url, url]) self.assertEqual(code, EXIT_OK) self.assertIn("\n\n", out) def test_fail_flag_ok_on_200(self): code, out, _ = _invoke(["--fail", self._url()]) self.assertEqual(code, EXIT_OK) self.assertNotIn("✗", out) def test_fail_flag_on_404(self): code, out, _ = _invoke(["--fail", self._url(self.nf_port)]) self.assertEqual(code, EXIT_HTTP) self.assertIn("✗", out) class TestCLIFailures(unittest.TestCase): def test_dns_failure(self): code, out, _ = _invoke(["http://no.such.host.invalid"]) self.assertEqual(code, EXIT_DNS) self.assertIn("FAILED", out) self.assertIn("✗", out) def test_connection_refused(self): port = _free_port() code, out, _ = _invoke([f"http://127.0.0.1:{port}"]) self.assertEqual(code, EXIT_CONNECT) self.assertIn("FAILED", out) def test_timeout(self): port = _black_hole_port() code, _, _ = _invoke([f"http://127.0.0.1:{port}", "--timeout", "200ms"]) self.assertEqual(code, EXIT_TIMEOUT) def test_worst_code_across_urls(self): server, port = _start_server( type("_H", (http.server.BaseHTTPRequestHandler,), { "do_GET": lambda self: (self.send_response(200), self.end_headers(), self.wfile.write(b"ok")), "log_message": lambda *a: None, }) ) try: code, _, _ = _invoke([ f"http://127.0.0.1:{port}", # exit 0 "http://no.such.host.invalid", # exit 2 ]) self.assertEqual(code, EXIT_DNS) finally: server.shutdown() def test_all_failed_single_header(self): code, out, _ = _invoke(["http://no.such.host.invalid"]) self.assertIn("(FAILED)", out) self.assertNotIn("0/1 succeeded", out) def test_all_failed_multi_header(self): code, out, _ = _invoke([ "-n", "3", "http://no.such.host.invalid", ]) self.assertIn("0/3 succeeded", out) def test_mixed_aggregate_shows_samples(self): port = _free_port() server, ok_port = _start_server(_OKHandler) try: code, out, _ = _invoke([ "-n", "1", f"http://127.0.0.1:{ok_port}", f"http://127.0.0.1:{port}", ]) self.assertEqual(code, EXIT_CONNECT) self.assertIn("200", out) finally: server.shutdown() class TestCLIJSON(unittest.TestCase): @classmethod def setUpClass(cls): cls.ok_server, cls.ok_port = _start_server(_OKHandler) cls.nf_server, cls.nf_port = _start_server(_NotFoundHandler) @classmethod def tearDownClass(cls): cls.ok_server.shutdown() cls.nf_server.shutdown() def _url(self, port=None): return f"http://127.0.0.1:{port or self.ok_port}" def test_json_success_schema(self): code, out, _ = _invoke(["--json", self._url()]) self.assertEqual(code, EXIT_OK) data = json.loads(out) self.assertIsInstance(data, list) self.assertEqual(len(data), 1) entry = data[0] self.assertEqual(entry["url"], self._url()) self.assertEqual(entry["status"], 200) self.assertEqual(entry["succeeded"], 1) self.assertEqual(entry["failed"], 0) self.assertIn("phases", entry) self.assertIn("total", entry["phases"]) self.assertNotIn("errors", entry) def test_json_phase_fields(self): code, out, _ = _invoke(["--json", self._url()]) data = json.loads(out) total = data[0]["phases"]["total"] for key in ("min_ms", "avg_ms", "max_ms"): self.assertIn(key, total) self.assertGreater(total[key], 0) def test_json_dns_failure_schema(self): code, out, _ = _invoke(["--json", "http://no.such.host.invalid"]) self.assertEqual(code, EXIT_DNS) data = json.loads(out) entry = data[0] self.assertEqual(entry["succeeded"], 0) self.assertEqual(entry["failed"], 1) self.assertNotIn("phases", entry) self.assertIn("errors", entry) self.assertEqual(entry["errors"][0]["phase"], "dns") def test_json_sampling(self): code, out, _ = _invoke(["--json", "-n", "3", self._url()]) self.assertEqual(code, EXIT_OK) data = json.loads(out) self.assertEqual(data[0]["succeeded"], 3) def test_json_multiple_urls_ordered(self): code, out, _ = _invoke([ "--json", self._url(), "http://no.such.host.invalid", ]) data = json.loads(out) self.assertEqual(len(data), 2) self.assertEqual(data[0]["status"], 200) self.assertEqual(data[1]["succeeded"], 0) def test_json_fail_flag_exit_code(self): code, out, _ = _invoke(["--json", "--fail", self._url(self.nf_port)]) self.assertEqual(code, EXIT_HTTP) data = json.loads(out) self.assertEqual(data[0]["status"], 404) def test_json_error_grouping(self): code, out, _ = _invoke([ "--json", "-n", "2", "http://no.such.host.invalid", ]) data = json.loads(out) self.assertEqual(data[0]["errors"][0]["count"], 2) class TestCLIVerbose(unittest.TestCase): @classmethod def setUpClass(cls): cls.ok_server, cls.ok_port = _start_server(_OKHandler) cls.nf_server, cls.nf_port = _start_server(_NotFoundHandler) @classmethod def tearDownClass(cls): cls.ok_server.shutdown() cls.nf_server.shutdown() def _url(self, port=None): return f"http://127.0.0.1:{port or self.ok_port}" def test_verbose_shows_ip(self): code, out, _ = _invoke(["--verbose", self._url()]) self.assertEqual(code, EXIT_OK) self.assertIn("IP", out) self.assertIn("127.0.0.1", out) def test_short_flag_v(self): code, out, _ = _invoke(["-v", self._url()]) self.assertEqual(code, EXIT_OK) self.assertIn("127.0.0.1", out) def test_verbose_no_tls_for_http(self): code, out, _ = _invoke(["--verbose", self._url()]) self.assertNotIn("TLS", out) self.assertNotIn("Cert", out) def test_verbose_shows_headers(self): code, out, _ = _invoke(["--verbose", self._url()]) # BaseHTTPServer sends Server and Content-Type self.assertIn("Server", out) def test_verbose_absent_without_flag(self): code, out, _ = _invoke([self._url()]) # Without --verbose, there is no IP label row (the URL contains the # IP but is not followed by a " IP" verbose block line) self.assertNotIn("\n IP ", out) def test_verbose_aggregate_shows_ip(self): code, out, _ = _invoke(["-v", "-n", "2", self._url()]) self.assertEqual(code, EXIT_OK) self.assertIn("127.0.0.1", out) self.assertIn("2 samples", out) def test_verbose_on_connect_fail_shows_ip(self): port = _free_port() code, out, _ = _invoke(["--verbose", f"http://127.0.0.1:{port}"]) self.assertEqual(code, EXIT_CONNECT) self.assertIn("FAILED", out) self.assertIn("127.0.0.1", out) def test_verbose_dns_fail_no_ip(self): code, out, _ = _invoke(["--verbose", "http://no.such.host.invalid"]) self.assertEqual(code, EXIT_DNS) # IP is unknown for DNS failures — verbose block is empty so no IP row self.assertNotIn("IP", out) def test_verbose_json_includes_ip(self): code, out, _ = _invoke(["--verbose", "--json", self._url()]) self.assertEqual(code, EXIT_OK) data = json.loads(out) self.assertIn("verbose", data[0]) self.assertEqual(data[0]["verbose"]["ip"], "127.0.0.1") def test_verbose_json_includes_headers(self): code, out, _ = _invoke(["--verbose", "--json", self._url()]) data = json.loads(out) self.assertIn("headers", data[0]["verbose"]) self.assertIsInstance(data[0]["verbose"]["headers"], dict) def test_no_verbose_key_in_json_without_flag(self): code, out, _ = _invoke(["--json", self._url()]) data = json.loads(out) self.assertNotIn("verbose", data[0]) def test_verbose_json_no_tls_for_http(self): code, out, _ = _invoke(["--verbose", "--json", self._url()]) data = json.loads(out) self.assertNotIn("tls_version", data[0]["verbose"]) self.assertNotIn("cert", data[0]["verbose"]) def test_verbose_json_connect_fail_has_ip(self): port = _free_port() code, out, _ = _invoke(["--verbose", "--json", f"http://127.0.0.1:{port}"]) data = json.loads(out) self.assertIn("verbose", data[0]) self.assertEqual(data[0]["verbose"]["ip"], "127.0.0.1") def test_verbose_json_dns_fail_no_verbose_object(self): code, out, _ = _invoke(["--verbose", "--json", "http://no.such.host.invalid"]) data = json.loads(out) # DNS failure: detail has no IP, so the verbose dict is empty → omitted self.assertNotIn("verbose", data[0])