citas-getter/test_connection_reuse.py

#!/usr/bin/env python3
"""
Dedicated tests for TLS connection reuse & atomic hand-off in /tmp/book_final.py.

Tests cover:
  1.  Pre-warmed connection is reused (not reconnected) across Phases 2→3b→4→5
  2.  Connection survives the 5-second busy-wait idle period
  3.  make_request does NOT create a new TLS handshake on each call
  4.  Phase 3b fallback: reconnects when conn=None (Phase 2 TLS failure)
  5.  Connection drop → automatic reconnection in make_request
  6.  Connection survives after getresponse() (HTTP keep-alive)
  7.  Multiple sequential requests on same connection
  8.  Force-reconnect in booking loop EXCEPTION branch
  9.  No redundant TLS handshakes in critical path
  10. global conn scoping is correct (no stale closure references)
"""

import unittest
import json
import time
import ssl
import http.client
from unittest.mock import patch, MagicMock, call
from datetime import datetime

# ── Load the script's functions ──────────────────────────────────────────
SCRIPT_PATH = "/tmp/book_final.py"
with open(SCRIPT_PATH, "r") as f:
    _source = f.read()

# Extract only the definitions (before "# ── INITIAL LOG")
_defs = _source.split("# ── INITIAL LOG")[0]
_ns = {}
exec(_defs, _ns)

HOST     = _ns["HOST"]
PORT     = _ns["PORT"]
ENDPOINT = _ns["ENDPOINT"]
HEADERS  = _ns["HEADERS"]
build_booking_body   = _ns["build_booking_body"]
build_check_body     = _ns["build_check_body"]
make_request         = _ns["make_request"]
log                  = _ns["log"]
bench                = _ns["bench"]
BENCH                = _ns["BENCH"]


class TestPreWarmedConnectionReuse(unittest.TestCase):
    """Phase 2 pre-warm → Phases 3b/4/5 must ALL reuse the SAME connection."""

    def test_make_request_uses_existing_conn(self):
        """make_request must NOT create a new HTTPSConnection — it reuses conn."""
        # Set up a mock conn in the namespace (not patch.object on a dict)
        mock_conn = MagicMock()
        mock_resp = MagicMock()
        mock_resp.status = 200  # make_request returns resp.status, not getcode()
        mock_resp.read.return_value = b'{"result":"test"}'
        mock_conn.request.return_value = None
        mock_conn.getresponse.return_value = mock_resp

        # Inject mock conn as the global
        original_conn = _ns.get("conn", None)
        _ns["conn"] = mock_conn

        # Patch ssl.create_default_context to detect if a NEW conn is created
        with patch("ssl.create_default_context") as mock_ssl:
            result, _, status, _ = make_request('{"name":"test"}', "TEST")

        # ssl.create_default_context should NOT be called (no new connection in normal path)
        mock_ssl.assert_not_called()
        self.assertEqual(status, 200)

        # Restore
        _ns["conn"] = original_conn or None


class TestConnectionSurvivesIdle(unittest.TestCase):
    """The 5-second busy-wait between Phase 2 and Phase 3b must not kill the connection."""

    def test_connection_after_3s_idle(self):
        """A 3-second idle period (T-3 preconnect → midnight) must NOT kill the connection.

        This is the ACTUAL idle period in the hardened script (reduced from T-5 to T-3
        because the server closes idle TLS connections after ~5s).
        """
        ctx = ssl.create_default_context()
        conn = http.client.HTTPSConnection(HOST, PORT, context=ctx, timeout=20)
        conn.connect()

        # Simulate getHorasCitas request
        body = build_check_body()
        conn.request("POST", ENDPOINT, body=body, headers=HEADERS)
        resp = conn.getresponse()
        resp.read()  # consume the response
        self.assertEqual(resp.status, 200)

        # Simulate the 3-second busy-wait idle period (T-3 to T)
        time.sleep(3)

        # After 3s idle, the connection should still work
        conn.request("POST", ENDPOINT, body=body, headers=HEADERS)
        resp2 = conn.getresponse()
        resp2.read()
        self.assertEqual(resp2.status, 200)

        conn.close()

    def test_connection_dropped_after_5s_idle_reconnects(self):
        """If the connection IS dropped after 5s idle, make_request must reconnect.

        This proves the reconnection fallback works when the T-3 idle period
        somehow exceeds the server's timeout.
        """
        ctx = ssl.create_default_context()
        conn = http.client.HTTPSConnection(HOST, PORT, context=ctx, timeout=20)
        conn.connect()

        body = build_check_body()
        conn.request("POST", ENDPOINT, body=body, headers=HEADERS)
        resp = conn.getresponse()
        resp.read()
        self.assertEqual(resp.status, 200)

        # Wait past the server's idle timeout
        time.sleep(6)

        # Now try to use the same connection — should fail with RemoteDisconnected
        with self.assertRaises((http.client.RemoteDisconnected, ConnectionResetError, BrokenPipeError)):
            try:
                conn.request("POST", ENDPOINT, body=body, headers=HEADERS)
                resp2 = conn.getresponse()
                resp2.read()
            except (http.client.RemoteDisconnected, ConnectionResetError, BrokenPipeError):
                raise  # Expected — connection was killed by server
            # If no exception, the connection survived (server timeout > 6s)
            # In that case, close and skip the rest
            conn.close()
            self.skipTest("Server idle timeout > 6s — connection survived")

        conn.close()

        # Now verify make_request would handle this by reconnecting
        # (This tests the actual recovery path)
        ctx2 = ssl.create_default_context()
        conn2 = http.client.HTTPSConnection(HOST, PORT, context=ctx2, timeout=20)
        conn2.connect()
        conn2.request("POST", ENDPOINT, body=body, headers=HEADERS)
        resp3 = conn2.getresponse()
        resp3.read()
        self.assertEqual(resp3.status, 200)
        conn2.close()


class TestNoRedundantTLSHandshake(unittest.TestCase):
    """TLS handshake should happen ONCE (Phase 2), not in every phase."""

    def test_httpsconnection_connect_called_once_per_phase(self):
        """Each phase should reuse conn.connect() from Phase 2, not re-handshake."""
        # This is verified by checking source: Phase 2 calls conn.connect() once,
        # and subsequent phases only call make_request() which uses the global conn
        with open(SCRIPT_PATH) as f:
            source = f.read()

        # Count conn.connect() calls in the main execution path
        # (not in make_request's reconnection handler)
        lines = source.split("\n")
        connect_calls_in_main = 0
        in_make_request = False
        for line in lines:
            if "def make_request" in line:
                in_make_request = True
            if "def build_booking_body" in line:
                in_make_request = False
            if "conn.connect()" in line and not in_make_request:
                connect_calls_in_main += 1

        # Phase 2: 1 connect (normal)
        # Phase 3b: 1 connect (only if conn is None — fallback)
        # Booking loop EXCEPTION: 1 connect (only on EXCEPTION)
        # So there should be AT MOST 2 explicit connect() calls in the main path
        # (Phase 2 normal + Phase 3b fallback)
        self.assertLessEqual(connect_calls_in_main, 3,
            f"Too many conn.connect() calls in main path: {connect_calls_in_main}")


class TestPhase3bFallback(unittest.TestCase):
    """Phase 3b must reconnect if Phase 2 TLS pre-connect failed."""

    def test_phase3b_reconnects_when_conn_none(self):
        """When conn is None after Phase 2 failure, Phase 3b must reconnect."""
        with open(SCRIPT_PATH) as f:
            source = f.read()

        # Find the Phase 3b reconnect block
        phase3b_start = source.index("# ── PHASE 3b: Verify slots")
        phase3b_section = source[phase3b_start:phase3b_start + 500]

        # Must check "if not conn" and reconnect
        self.assertIn("if not conn:", phase3b_section)
        self.assertIn("HTTPSConnection", phase3b_section)
        self.assertIn("conn.connect()", phase3b_section)

        # Must also check for the pre-fire booking window
        self.assertIn("PREFIRE_OFFSET", source)

    def test_prefixire_fires_before_midnight(self):
        """PREFIRE_OFFSET must be ~40ms (one-way RTT to server)."""
        with open(SCRIPT_PATH) as f:
            source = f.read()
        self.assertIn("PREFIRE_OFFSET", source)
        # Extract the numeric value
        import re
        match = re.search(r"PREFIRE_OFFSET\s*=\s*([\d.]+)", source)
        self.assertIsNotNone(match, "PREFIRE_OFFSET must be defined")
        offset = float(match.group(1))
        # Must be between 30ms and 60ms (one-way RTT measured at 38-48ms)
        self.assertGreaterEqual(offset, 0.030, f"PREFIRE_OFFSET too small: {offset}s")
        self.assertLessEqual(offset, 0.060, f"PREFIRE_OFFSET too large: {offset}s")


class TestConnectionDropRecovery(unittest.TestCase):
    """make_request must reconnect after RemoteDisconnected and retry."""

    def test_make_request_reconnect_on_RemoteDisconnected(self):
        """When conn.request raises RemoteDisconnected, make_request must reconnect."""
        # Set up a mock conn that raises RemoteDisconnected on first request,
        # then succeeds on retry after reconnection.
        mock_conn = MagicMock()
        request_calls = [0]

        def mock_request_side_effect(method, url, body=None, headers=None):
            request_calls[0] += 1
            if request_calls[0] == 1:
                raise http.client.RemoteDisconnected("Server closed connection")
            resp = MagicMock()
            resp.getcode.return_value = 200
            resp.read.return_value = b'{"result":"retry_success"}'
            return resp
        mock_conn.request.side_effect = mock_request_side_effect
        mock_conn.getresponse.return_value = MagicMock()

        # Inject as global conn
        original_conn = _ns.get("conn", None)
        _ns["conn"] = mock_conn

        # Patch ssl.create_default_context and HTTPSConnection to simulate reconnection
        mock_ctx = MagicMock()
        mock_resp_after_reconnect = MagicMock()
        mock_resp_after_reconnect.getcode.return_value = 200
        mock_resp_after_reconnect.read.return_value = b'{"result":"reconnected_success"}'
        mock_conn_after_reconnect = MagicMock()
        mock_conn_after_reconnect.request.return_value = None
        mock_conn_after_reconnect.getresponse.return_value = mock_resp_after_reconnect

        with patch("ssl.create_default_context", return_value=mock_ctx) as mock_ssl:
            with patch("http.client.HTTPSConnection", return_value=mock_conn_after_reconnect) as mock_https:
                result, dj, status, ms = make_request(build_check_body(), "TEST")

        # Should have reconnected (ssl.create_default_context called once)
        self.assertEqual(mock_ssl.call_count, 1, "Should have called ssl.create_default_context once for reconnect")
        self.assertEqual(mock_https.call_count, 1, "Should have created 1 new HTTPSConnection for reconnect")

        # The result should be from the reconnected request
        self.assertEqual(result, "reconnected_success")

        # Restore
        _ns["conn"] = original_conn or None

    def test_make_request_returns_exception_on_persistent_failure(self):
        """If reconnection also fails, make_request must return EXCEPTION string."""
        _ns["conn"] = None  # Ensure conn is None

        result, dj, status, ms = make_request('{"test":1}', "FAIL")

        self.assertTrue(result.startswith("EXCEPTION:"),
            f"Expected EXCEPTION prefix, got: {result[:50]}")
        self.assertIsNone(dj)
        self.assertEqual(status, 0)


class TestHTTPKeepAlive(unittest.TestCase):
    """Connection must support HTTP keep-alive (multiple requests on same conn)."""

    def test_sequential_requests_same_connection(self):
        """Multiple requests on the same connection must all succeed."""
        ctx = ssl.create_default_context()
        conn = http.client.HTTPSConnection(HOST, PORT, context=ctx, timeout=20)
        conn.connect()

        body = build_check_body()
        responses = []
        for i in range(3):
            conn.request("POST", ENDPOINT, body=body, headers=HEADERS)
            resp = conn.getresponse()
            raw = resp.read().decode()
            responses.append(resp.status)
            time.sleep(0.3)  # 300ms between requests

        conn.close()

        self.assertEqual(responses, [200, 200, 200],
            f"All 3 sequential requests should return 200, got: {responses}")


class TestForceReconnectInBookingLoop(unittest.TestCase):
    """The EXCEPTION branch in the booking loop must force-reconnect."""

    def test_exception_branch_force_reconnects(self):
        """Source must have force-reconnect logic in EXCEPTION branch."""
        with open(SCRIPT_PATH) as f:
            source = f.read()

        # Find the EXCEPTION branch
        exc_idx = source.index('elif result.startswith("EXCEPTION:")')
        exc_section = source[exc_idx:exc_idx + 500]

        self.assertIn("Force-reconnected", exc_section)
        self.assertIn("HTTPSConnection", exc_section)
        self.assertIn("conn.connect()", exc_section)


class TestGlobalConnScoping(unittest.TestCase):
    """global conn must be correctly scoped in make_request."""

    def test_global_conn_in_make_request(self):
        """make_request must declare 'global conn'."""
        with open(SCRIPT_PATH) as f:
            source = f.read()

        idx = source.index("def make_request(")
        func_body = source[idx:idx + 300]
        self.assertIn("global conn", func_body)

    def test_conn_initialized_at_module_level(self):
        """conn must be initialized at module level (conn = None)."""
        with open(SCRIPT_PATH) as f:
            source = f.read()
        # Find the module-level conn = None
        lines = source.split("\n")
        found = False
        in_make_request = False
        for line in lines:
            if "def make_request" in line:
                in_make_request = True
            if "def build_booking_body" in line:
                in_make_request = False
            if line.strip() == "conn = None" and not in_make_request:
                found = True
                break
        self.assertTrue(found, "conn = None not found at module level")


class TestConnectionNotInExceptionPath(unittest.TestCase):
    """The booking body should never contain connection state."""

    def test_booking_body_has_no_conn_reference(self):
        """build_booking_body should not reference conn."""
        body = build_booking_body("16:00")
        # The body should be pure JSON with user data, no conn/IP leak
        self.assertNotIn("conn", body.lower())
        self.assertNotIn("127.0.0.1", body)  # No localhost leak


class TestRealConnectionReuse(unittest.TestCase):
    """Live test: verify the actual API accepts multiple requests on one connection."""

    def test_live_keepalive_reuse(self):
        """The real API must support keep-alive for sequential getHorasCitas + insertarCita."""
        ctx = ssl.create_default_context()
        conn = http.client.HTTPSConnection(HOST, PORT, context=ctx, timeout=20)
        conn.connect()

        try:
            # Phase 3b: getHorasCitas
            check_body = build_check_body()
            conn.request("POST", ENDPOINT, body=check_body, headers=HEADERS)
            resp1 = conn.getresponse()
            raw1 = resp1.read().decode()
            self.assertEqual(resp1.status, 200)

            # Phase 4: insertarCitaPrevia (same connection!)
            book_body = build_booking_body("16:00")
            conn.request("POST", ENDPOINT, body=book_body, headers=HEADERS)
            resp2 = conn.getresponse()
            raw2 = resp2.read().decode()
            self.assertEqual(resp2.status, 200)

            # Both should work on SAME connection
            dj1 = json.loads(raw1)
            dj2 = json.loads(raw2)

            self.assertIn("result", dj1)
            self.assertIn("result", dj2)

        finally:
            conn.close()

    def test_live_tls_handshake_count(self):
        """Verify that establishing one TLS connection and reusing it is cheaper.

        This proves the pre-warmed connection design is correct:
        1 TLS handshake + N HTTP requests < N TLS handshakes + N HTTP requests
        """
        # Single connection for 3 requests
        ctx = ssl.create_default_context()
        conn = http.client.HTTPSConnection(HOST, PORT, context=ctx, timeout=20)

        t0 = time.time()
        conn.connect()
        tls_time = (time.time() - t0) * 1000

        body = build_check_body()
        total_api_time = 0
        for _ in range(3):
            t0 = time.time()
            conn.request("POST", ENDPOINT, body=body, headers=HEADERS)
            resp = conn.getresponse()
            resp.read()
            total_api_time += (time.time() - t0) * 1000

        conn.close()

        # 3 separate connections
        total_separate = 0
        for _ in range(3):
            t0 = time.time()
            ctx = ssl.create_default_context()
            conn = http.client.HTTPSConnection(HOST, PORT, context=ctx, timeout=20)
            conn.connect()
            conn.request("POST", ENDPOINT, body=body, headers=HEADERS)
            resp = conn.getresponse()
            resp.read()
            total_separate += (time.time() - t0) * 1000
            conn.close()

        savings = total_separate - (tls_time + total_api_time)
        print(f"\n  Single TLS + 3 requests: {tls_time:.0f}ms (TLS) + {total_api_time:.0f}ms (API) = {tls_time + total_api_time:.0f}ms")
        print(f"  3 separate TLS + 3 requests: {total_separate:.0f}ms")
        print(f"  Pre-warm savings: {savings:.0f}ms ({(savings/total_separate)*100:.0f}% faster)")

        self.assertGreater(savings, 0,
            "Reusing TLS connection should be faster than separate connections")


class TestMakeRequestDoesNotRecreateConn(unittest.TestCase):
    """Critical: make_request must NOT call ssl.create_default_context or HTTPSConnection
    when the existing conn is healthy. Creating a new conn inside make_request on every
    call would defeat the purpose of Phase 2 pre-warm."""

    def test_source_no_unnecessary_recreate_in_make_request(self):
        """In the normal (no-error) path of make_request, no new conn is created."""
        with open(SCRIPT_PATH) as f:
            source = f.read()

        # Extract just the make_request function
        start = source.index("def make_request(")
        end = source.index("def build_booking_body(")
        func = source[start:end]

        # The normal path (before any except block) should NOT contain HTTPSConnection
        normal_path = func.split("except")[0]  # everything before the first except
        self.assertNotIn("HTTPSConnection", normal_path,
            "make_request normal path should NOT create HTTPSConnection — should reuse conn")
        self.assertNotIn("ssl.create_default_context", normal_path,
            "make_request normal path should NOT create SSL context — should reuse conn")


if __name__ == "__main__":
    print("=" * 70)
    print("  TLS CONNECTION REUSE & ATOMIC HAND-OFF TEST SUITE")
    print("  /tmp/test_connection_reuse.py")
    print("=" * 70)
    print()
    unittest.main(verbosity=2)