atlas-iac/testing/tests/test_hermes_deadline_stream.py

"""Watchdog contracts for absolute deadlines on streaming Git exchanges."""

from __future__ import annotations

import threading
import time
import urllib.request
from http.server import BaseHTTPRequestHandler, HTTPServer

import pytest

from testing.tests.test_hermes_scm_broker_support import Response, _load


def test_stream_deadline_validates_bounds_and_counts_down():
    module = _load("deadline_http")
    for timeout in (0, -1, module.MAX_STREAM_SECONDS + 1):
        with pytest.raises(module.PolicyError, match="bounds are invalid"):
            module.StreamDeadline(timeout)
    guard = module.StreamDeadline(30)
    try:
        first = guard.remaining()
        assert 0 < first <= 30
    finally:
        guard.cancel()


class _FakeTimer:
    def __init__(self, timeout, callback):
        self.timeout = timeout
        self.callback = callback
        self.daemon = False
        self.started = False
        self.cancelled = False

    def start(self):
        self.started = True

    def cancel(self):
        self.cancelled = True


class _FakeSocket:
    def __init__(self, fail: bool = False):
        self.fail = fail
        self.shut = False

    def shutdown(self, _how):
        if self.fail:
            raise OSError("already gone")
        self.shut = True


class _FakeConnection:
    def __init__(self, sock=None):
        self.sock = sock
        self.closed = False

    def close(self):
        self.closed = True


def test_expiry_shuts_down_and_closes_every_tracked_connection():
    module = _load("deadline_http")
    guard = module.StreamDeadline(30, timer=_FakeTimer)
    assert guard._timer.started and guard._timer.daemon
    healthy = _FakeConnection(_FakeSocket())
    broken = _FakeConnection(_FakeSocket(fail=True))
    bare = _FakeConnection()
    for connection in (healthy, broken, bare):
        assert guard._track(connection) is connection
    guard._expire()
    assert guard.expired
    assert healthy.sock.shut and healthy.closed
    assert broken.closed and not broken.sock.shut
    assert bare.closed
    with pytest.raises(module.PolicyError, match="deadline exceeded"):
        guard.remaining()
    late = _FakeConnection()
    with pytest.raises(module.PolicyError, match="deadline exceeded"):
        guard._track(late)
    assert late.closed
    guard.cancel()
    assert guard._timer.cancelled


def test_handlers_build_guarded_connections_without_network(monkeypatch):
    module = _load("deadline_http")
    guard = module.StreamDeadline(30, timer=_FakeTimer)
    http_handler, https_handler = guard.handlers()
    assert isinstance(http_handler, urllib.request.HTTPHandler)
    assert isinstance(https_handler, urllib.request.HTTPSHandler)
    for handler, opener_name in (
        (http_handler, "http_open"),
        (https_handler, "https_open"),
    ):
        seen = {}
        monkeypatch.setattr(
            handler,
            "do_open",
            lambda factory, req, seen=seen: seen.update(factory=factory, req=req),
        )
        getattr(handler, opener_name)("request")
        connection = seen["factory"]("127.0.0.1", timeout=1)
        assert connection in guard._connections
        connection.close()
    guard.cancel()


class _TrickleHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        time.sleep(10)

    def log_message(self, *_args):
        return


def test_watchdog_interrupts_a_hung_streaming_exchange():
    module = _load("deadline_http")
    server = HTTPServer(("127.0.0.1", 0), _TrickleHandler)
    threading.Thread(target=server.serve_forever, daemon=True).start()
    guard = module.StreamDeadline(1.0)
    opener = urllib.request.build_opener(*guard.handlers())
    started = time.monotonic()
    try:
        with pytest.raises(OSError):
            opener.open(
                f"http://127.0.0.1:{server.server_port}/slow", timeout=30
            ).read()
        assert time.monotonic() - started < 5
        assert guard.expired
    finally:
        guard.cancel()
        server.shutdown()


def test_upstream_git_request_runs_under_one_wall_clock_budget(monkeypatch):
    broker = _load("scm_broker")
    seen = {}

    class RecordingGuard:
        def __init__(self, timeout, **_kwargs):
            seen["timeout"] = timeout
            self.cancelled = False
            seen["guard"] = self

        def remaining(self):
            return 12.5

        def cancel(self):
            self.cancelled = True

        def handlers(self):
            return ()

    monkeypatch.setattr(broker.deadline_http, "StreamDeadline", RecordingGuard)

    def opener(request, timeout):
        seen["opener_timeout"] = timeout
        return Response(b"result", content_type="application/x-git-upload-pack-result")

    result = broker._upstream_git_request(
        "/atlas/cassandra.git/git-upload-pack",
        method="POST",
        body=b"request",
        content_type="application/x-git-upload-pack-request",
        expected_type="application/x-git-upload-pack-result",
        token="sentinel",
        opener=opener,
    )
    assert result == b"result"
    assert seen["timeout"] == broker.UPSTREAM_DEADLINE_SECONDS
    assert seen["opener_timeout"] == 12.5
    assert seen["guard"].cancelled


def test_upstream_git_request_cancels_guard_when_deadline_already_passed(monkeypatch):
    broker = _load("scm_broker")
    cancelled = []

    class ExpiredGuard:
        def __init__(self, _timeout, **_kwargs):
            pass

        def remaining(self):
            raise broker.PolicyError("HTTP stream deadline exceeded")

        def cancel(self):
            cancelled.append(True)

        def handlers(self):
            return ()

    monkeypatch.setattr(broker.deadline_http, "StreamDeadline", ExpiredGuard)
    with pytest.raises(broker.PolicyError, match="deadline exceeded"):
        broker._upstream_git_request(
            "/atlas/cassandra.git/git-upload-pack",
            method="POST",
            body=b"request",
            content_type="application/x-git-upload-pack-request",
            expected_type="application/x-git-upload-pack-result",
            token="sentinel",
            opener=lambda *_a, **_k: None,
        )
    assert cancelled == [True]


def test_guarded_opener_wires_redirect_rejection_and_deadline_handlers():
    broker = _load("scm_broker")
    module = _load("deadline_http")
    guard = module.StreamDeadline(30, timer=_FakeTimer)
    opener = broker._guarded_opener(guard)
    assert callable(opener)
    handlers = opener.__self__.handlers
    assert any(isinstance(item, broker.RejectRedirect) for item in handlers)
    assert any(isinstance(item, urllib.request.HTTPSHandler) for item in handlers)
    guard.cancel()