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

import http.server
import socket
import threading
import time
import unittest

from latprobe.probe import Options, Phase, Result, VerboseDetail, measure


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:
    """Bind a port that accepts TCP but never sends any data (triggers TTFB timeout)."""
    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


class TestMeasureSuccess(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 test_all_phases_present_on_success(self):
        r = measure(f"http://127.0.0.1:{self.ok_port}")
        self.assertIsNone(r.err)
        self.assertEqual(r.fail_phase, "")
        self.assertEqual(r.status_code, 200)
        self.assertTrue(r.dns.present)
        self.assertTrue(r.connect.present)
        self.assertFalse(r.tls.present)   # http — no TLS
        self.assertTrue(r.ttfb.present)
        self.assertTrue(r.transfer.present)
        self.assertTrue(r.total.present)

    def test_total_gte_sum_of_phases(self):
        r = measure(f"http://127.0.0.1:{self.ok_port}")
        phase_sum = r.dns.ms + r.connect.ms + r.ttfb.ms + r.transfer.ms
        self.assertGreaterEqual(r.total.ms, phase_sum * 0.9)

    def test_status_code_404(self):
        r = measure(f"http://127.0.0.1:{self.nf_port}")
        self.assertIsNone(r.err)
        self.assertEqual(r.status_code, 404)
        self.assertTrue(r.total.present)

    def test_options_default_timeout(self):
        opts = Options()
        r = measure(f"http://127.0.0.1:{self.ok_port}", opts)
        self.assertIsNone(r.err)
        self.assertEqual(r.status_code, 200)


class TestMeasureFailures(unittest.TestCase):

    def test_dns_failure(self):
        r = measure("http://no.such.host.invalid")
        self.assertEqual(r.fail_phase, "dns")
        self.assertIsNotNone(r.err)
        self.assertFalse(r.dns.present)
        self.assertFalse(r.connect.present)
        self.assertTrue(r.total.present)

    def test_connection_refused(self):
        port = _free_port()
        r = measure(f"http://127.0.0.1:{port}")
        self.assertEqual(r.fail_phase, "connect")
        self.assertIsNotNone(r.err)
        self.assertTrue(r.dns.present)
        self.assertTrue(r.connect.present)
        self.assertFalse(r.ttfb.present)
        self.assertTrue(r.total.present)

    def test_ttfb_timeout(self):
        port = _black_hole_port()
        r = measure(f"http://127.0.0.1:{port}", Options(timeout=0.2))
        self.assertEqual(r.fail_phase, "timeout")
        self.assertIsNotNone(r.err)
        self.assertTrue(r.dns.present)
        self.assertTrue(r.connect.present)
        self.assertTrue(r.total.present)

    def test_bad_scheme(self):
        r = measure("ftp://example.com")
        self.assertEqual(r.fail_phase, "request")
        self.assertIsNotNone(r.err)
        self.assertFalse(r.total.present)

    def test_partial_phases_preserved_on_connect_fail(self):
        port = _free_port()
        r = measure(f"http://127.0.0.1:{port}")
        self.assertTrue(r.dns.present, "dns should be recorded before connect")
        self.assertTrue(r.connect.present, "connect duration recorded even on refusal")
        self.assertGreater(r.dns.ms, 0)
        self.assertGreater(r.connect.ms, 0)


class TestVerboseDetail(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.ok_server, cls.ok_port = _start_server(_OKHandler)

    @classmethod
    def tearDownClass(cls):
        cls.ok_server.shutdown()

    def _url(self):
        return f"http://127.0.0.1:{self.ok_port}"

    def test_detail_none_without_verbose(self):
        r = measure(self._url())
        self.assertIsNone(r.detail)

    def test_detail_present_with_verbose(self):
        r = measure(self._url(), Options(verbose=True))
        self.assertIsInstance(r.detail, VerboseDetail)

    def test_resolved_ip_set_on_success(self):
        r = measure(self._url(), Options(verbose=True))
        self.assertEqual(r.detail.resolved_ip, "127.0.0.1")

    def test_no_tls_fields_for_http(self):
        r = measure(self._url(), Options(verbose=True))
        self.assertEqual(r.detail.tls_version, "")
        self.assertEqual(r.detail.tls_cipher, "")
        self.assertEqual(r.detail.tls_bits, 0)
        self.assertIsNone(r.detail.cert)

    def test_headers_populated_on_success(self):
        r = measure(self._url(), Options(verbose=True))
        self.assertIsInstance(r.detail.headers, dict)
        # BaseHTTPServer always sends Content-Type for 200 responses
        self.assertTrue(len(r.detail.headers) > 0)

    def test_ip_set_on_connect_fail(self):
        port = _free_port()
        r = measure(f"http://127.0.0.1:{port}", Options(verbose=True))
        self.assertEqual(r.fail_phase, "connect")
        self.assertEqual(r.detail.resolved_ip, "127.0.0.1")
        self.assertEqual(r.detail.headers, {})

    def test_detail_none_on_dns_failure_without_verbose(self):
        r = measure("http://no.such.host.invalid")
        self.assertIsNone(r.detail)

    def test_detail_empty_ip_on_dns_failure(self):
        r = measure("http://no.such.host.invalid", Options(verbose=True))
        self.assertIsNotNone(r.detail)
        self.assertEqual(r.detail.resolved_ip, "")

    def test_headers_not_populated_on_connect_fail(self):
        port = _free_port()
        r = measure(f"http://127.0.0.1:{port}", Options(verbose=True))
        self.assertEqual(r.detail.headers, {})