atlas-iac/testing/tests/test_hermes_scm_broker_streaming.py

"""Bounded I/O and concurrency contracts for the Hermes SCM broker."""

from __future__ import annotations

import io
import json
from email.message import Message

import pytest

from testing.tests.test_hermes_scm_broker_support import Response, _load


def test_agent_broker_client_sends_no_credential_or_authorization_header():
    client = _load("scm_broker_client")
    seen = []

    def opener(request, timeout):
        seen.append((request, timeout))
        return Response(b"{}", content_type="application/json")

    assert client.read("/api/v1/repos/atlas/cassandra", opener=opener) == b"{}"
    request = seen[0][0]
    assert request.full_url == client.BROKER_ORIGIN + "/v1/metadata"
    assert request.get_header("Authorization") is None
    assert json.loads(request.data) == {"path": "/api/v1/repos/atlas/cassandra"}


@pytest.mark.parametrize("raw", ["", "01", "9" * 4000, "134217729", "１２"])
def test_broker_content_length_is_canonical_and_bounded(raw: str):
    broker = _load("scm_broker")
    headers = Message()
    headers["Content-Length"] = raw

    with pytest.raises(broker.PolicyError, match="request length"):
        broker._content_length(headers, broker.MAX_GIT_REQUEST)


def test_large_git_exchange_rolls_to_disk_and_scans_chunk_boundaries():
    broker = _load("scm_broker")
    token = "runtime-sentinel"
    safe = b"x" * (broker.SPOOL_MEMORY_LIMIT + 1)
    inbound_timeouts = []
    spool, length = broker._spool_bounded(
        io.BytesIO(safe),
        broker.MAX_GIT_REQUEST,
        len(safe),
        token=token,
        context="Git request",
        deadline=broker.time.monotonic() + 30,
        set_timeout=inbound_timeouts.append,
    )
    try:
        assert length == len(safe)
        assert spool._rolled is True
        assert spool.read(4) == b"xxxx"
        assert inbound_timeouts
        assert all(0 < value <= 30 for value in inbound_timeouts)
    finally:
        spool.close()

    crossing = b"x" * (broker.STREAM_CHUNK - 5) + token.encode() + b"tail"
    with pytest.raises(broker.PolicyError, match="credential material"):
        broker._spool_bounded(
            io.BytesIO(crossing),
            broker.MAX_GIT_REQUEST,
            len(crossing),
            token=token,
            context="Git request",
            deadline=broker.time.monotonic() + 30,
        )


def test_large_upstream_response_is_streamed_through_bounded_spool():
    broker = _load("scm_broker")
    body = b"z" * (broker.SPOOL_MEMORY_LIMIT + 1)
    seen = []

    def opener(request, **_kwargs):
        seen.append(request)
        assert request.get_header("Content-length") == "7"
        assert request.data.read() == b"request"
        return Response(body, content_type="application/x-git-upload-pack-result")

    streamed = broker._upstream_git_request(
        "/atlas/cassandra.git/git-upload-pack",
        method="POST",
        body=io.BytesIO(b"request"),
        body_length=7,
        content_type="application/x-git-upload-pack-request",
        expected_type="application/x-git-upload-pack-result",
        token="runtime-sentinel",
        stream_result=True,
        opener=opener,
    )
    spool, length = streamed
    try:
        assert length == len(body)
        assert spool._rolled is True
        assert spool.read(3) == b"zzz"
        assert len(seen) == 1
    finally:
        spool.close()


def test_body_deadline_fails_closed_before_reading():
    broker = _load("scm_broker")
    with pytest.raises(broker.PolicyError, match="deadline"):
        broker._spool_bounded(
            io.BytesIO(b"body"),
            broker.MAX_GIT_REQUEST,
            4,
            token="runtime-sentinel",
            context="Git request",
            deadline=broker.time.monotonic() - 1,
        )


def test_upstream_stream_applies_absolute_socket_deadline():
    broker = _load("scm_broker")

    class DeadlineSocket:
        def __init__(self):
            self.values = []

        def settimeout(self, value):
            self.values.append(value)

    response = Response(b"safe", content_type="application/x-git-upload-pack-result")
    sock = DeadlineSocket()
    response.fp = type("FP", (), {"raw": type("Raw", (), {"_sock": sock})()})()
    spool, length = broker._spool_response(
        response, broker.MAX_GIT_RESPONSE, "sentinel", broker.INBOUND_BODY_TIMEOUT
    )
    try:
        assert length == 4
        assert sock.values
        assert all(0 < value <= broker.INBOUND_BODY_TIMEOUT for value in sock.values)
    finally:
        spool.close()


def test_broker_rejects_work_when_concurrency_slots_are_exhausted():
    broker = _load("scm_broker")

    class FakeSocket:
        def __init__(self):
            self.value = b""

        def sendall(self, value):
            self.value += value

        def shutdown(self, _how):
            return None

        def close(self):
            return None

    server = broker.BoundedThreadingHTTPServer(("127.0.0.1", 0), broker.BrokerHandler)
    try:
        assert broker.BoundedThreadingHTTPServer.__module__ == "scm_broker_server"
        assert server._slots.acquire(blocking=False)
        assert server._slots.acquire(blocking=False)
        assert server._slots.acquire(blocking=False)
        request = FakeSocket()
        server.process_request(request, ("127.0.0.1", 1))
        assert b"503 Service Unavailable" in request.value
    finally:
        server._slots.release()
        server._slots.release()
        server._slots.release()
        server.server_close()


def test_three_simultaneous_broker_calls_have_bounded_capacity():
    server_module = _load("scm_broker_server")
    semaphore = server_module.threading.BoundedSemaphore(
        server_module.MAX_CONCURRENT_REQUESTS
    )
    assert server_module.MAX_CONCURRENT_REQUESTS >= 3
    assert [semaphore.acquire(blocking=False) for _ in range(3)] == [True, True, True]
    assert semaphore.acquire(blocking=False) is False


def test_header_deadline_is_absolute_against_slow_trickle(monkeypatch):
    server_module = _load("scm_broker_server")

    class Stream:
        def __init__(self):
            self.value = io.BytesIO(b"GET /healthz HTTP/1.1\r\n")

        def read(self, length):
            return self.value.read(length)

    class Connection:
        def __init__(self):
            self.timeouts = []

        def settimeout(self, value):
            self.timeouts.append(value)

    ticks = iter((0.0, 0.02, 0.04, 0.06, 0.08, 0.11))
    monkeypatch.setattr(server_module.time, "monotonic", lambda: next(ticks))
    connection = Connection()
    reader = server_module._AbsoluteDeadlineReader(Stream(), connection)
    reader.begin(0.1)

    with pytest.raises(TimeoutError, match="absolute SCM header deadline"):
        reader.readline(65537)

    assert connection.timeouts == pytest.approx([0.08, 0.06, 0.04, 0.02])


def test_broker_bounds_header_count_size_and_duplicate_lengths():
    broker = _load("scm_broker")

    def validate(headers):
        handler = object.__new__(broker.BrokerHandler)
        handler.headers = headers
        handler._validate_headers()

    too_many = Message()
    for index in range(broker.MAX_HEADERS + 1):
        too_many[f"X-Test-{index}"] = "safe"
    with pytest.raises(broker.PolicyError, match="too many headers"):
        validate(too_many)

    too_large = Message()
    too_large["X-Test"] = "x" * broker.MAX_HEADER_BYTES
    with pytest.raises(broker.PolicyError, match="safe limit"):
        validate(too_large)

    duplicate = Message()
    duplicate["Content-Length"] = "1"
    duplicate["Content-Length"] = "1"
    with pytest.raises(broker.PolicyError, match="duplicate Content-Length"):
        validate(duplicate)