http-latency-prober/python/tests/test_cli.py

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])