cleveragents-core/features/steps/asgi_app_steps.py

"""Step definitions for ASGI protocol behavior scenarios."""

from __future__ import annotations

import asyncio
import logging
from collections import deque
from typing import Any

from behave import given, then, when
from behave.runner import Context

from cleveragents.a2a.asgi import app

SendMessage = dict[str, Any]


@given("the ASGI app module is loaded")
def step_asgi_module_loaded(context: Context) -> None:
    context.asgi_app = app


@when('I send an HTTP GET request to "{path}" through the ASGI app')
def step_send_http_request(context: Context, path: str) -> None:
    sent_messages: list[SendMessage] = []

    async def receive() -> dict[str, Any]:
        return {"type": "http.request", "body": b"", "more_body": False}

    async def send(message: SendMessage) -> None:
        sent_messages.append(message)

    scope = {"type": "http", "method": "GET", "path": path}
    asyncio.run(context.asgi_app(scope, receive, send))
    context.asgi_sent_messages = sent_messages


@when('I send an HTTP POST request to "{path}" through the ASGI app')
def step_send_http_post_request(context: Context, path: str) -> None:
    sent_messages: list[SendMessage] = []

    async def receive() -> dict[str, Any]:
        return {"type": "http.request", "body": b"", "more_body": False}

    async def send(message: SendMessage) -> None:
        sent_messages.append(message)

    scope = {"type": "http", "method": "POST", "path": path}
    asyncio.run(context.asgi_app(scope, receive, send))
    context.asgi_sent_messages = sent_messages


@then("the HTTP response status should be {status:d}")
def step_http_status(context: Context, status: int) -> None:
    start = next(
        msg
        for msg in context.asgi_sent_messages
        if msg.get("type") == "http.response.start"
    )
    assert start.get("status") == status, (
        f"Expected HTTP status {status}, got {start.get('status')}"
    )


@then('the HTTP response body should be "{body}"')
def step_http_body(context: Context, body: str) -> None:
    response_body = next(
        msg
        for msg in context.asgi_sent_messages
        if msg.get("type") == "http.response.body"
    )
    normalized_body = body.replace(r"\"", '"')
    assert response_body.get("body") == normalized_body.encode("utf-8"), (
        f"Expected body {normalized_body!r}, got {response_body.get('body')!r}"
    )


@when("I run ASGI lifespan startup then shutdown")
def step_run_lifespan(context: Context) -> None:
    sent_messages: list[SendMessage] = []
    incoming = deque(
        [
            {"type": "lifespan.startup"},
            {"type": "lifespan.shutdown"},
        ]
    )

    async def receive() -> dict[str, Any]:
        if not incoming:
            raise AssertionError(
                "ASGI app issued an unexpected extra receive() call; "
                "all queued lifespan messages have already been consumed"
            )
        return incoming.popleft()

    async def send(message: SendMessage) -> None:
        sent_messages.append(message)

    scope = {"type": "lifespan"}
    asyncio.run(context.asgi_app(scope, receive, send))
    context.asgi_sent_messages = sent_messages


@then("the ASGI app should emit lifespan completion messages")
def step_assert_lifespan_messages(context: Context) -> None:
    message_types = [str(msg.get("type")) for msg in context.asgi_sent_messages]
    assert message_types == [
        "lifespan.startup.complete",
        "lifespan.shutdown.complete",
    ], f"Unexpected lifespan messages: {message_types}"


@then("no HTTP response frames should be emitted")
def step_no_http_frames(context: Context) -> None:
    message_types = [str(msg.get("type")) for msg in context.asgi_sent_messages]
    assert not any(
        msg_type.startswith("http.response") for msg_type in message_types
    ), f"Unexpected HTTP response frames: {message_types}"


@when("I invoke the ASGI app with a websocket scope")
def step_websocket_scope(context: Context) -> None:
    sent_messages: list[SendMessage] = []

    async def receive() -> dict[str, Any]:
        return {"type": "websocket.connect"}

    async def send(message: SendMessage) -> None:
        sent_messages.append(message)

    scope = {"type": "websocket", "path": "/ws"}
    asyncio.run(context.asgi_app(scope, receive, send))
    context.asgi_sent_messages = sent_messages


@when("I invoke the ASGI app with an unsupported scope type")
def step_unsupported_scope(context: Context) -> None:
    async def receive() -> dict[str, Any]:
        return {"type": "unsupported.receive"}

    async def send(message: SendMessage) -> None:
        del message

    scope = {"type": "unsupported", "path": "/"}
    try:
        asyncio.run(context.asgi_app(scope, receive, send))
    except RuntimeError as exc:
        context.asgi_error = exc
        return

    raise AssertionError("Expected RuntimeError for unsupported ASGI scope type")


@then("the ASGI app should emit a websocket close frame")
def step_websocket_close_frame(context: Context) -> None:
    assert context.asgi_sent_messages == [{"type": "websocket.close", "code": 1008}], (
        f"Unexpected websocket frames: {context.asgi_sent_messages}"
    )


@then("the ASGI invocation should raise an unsupported scope runtime error")
def step_unsupported_scope_error(context: Context) -> None:
    error = getattr(context, "asgi_error", None)
    assert isinstance(error, RuntimeError), f"Expected RuntimeError, got {error!r}"
    assert "Unsupported ASGI scope type" in str(error), (
        f"Unexpected runtime error message: {error}"
    )


# --- Allow header assertion for 405 responses ---


@then('the HTTP response should include an Allow header with value "{value}"')
def step_http_allow_header(context: Context, value: str) -> None:
    start = next(
        msg
        for msg in context.asgi_sent_messages
        if msg.get("type") == "http.response.start"
    )
    headers: list[tuple[bytes, bytes]] = start.get("headers", [])
    allow_values = [v.decode("utf-8") for k, v in headers if k.lower() == b"allow"]
    assert allow_values, f"Expected Allow header, but none found in: {headers}"
    assert value in allow_values, (
        f"Expected Allow header value {value!r}, got {allow_values!r}"
    )


# --- Security-hardening header assertions ---


@then("the HTTP response should include a content-length header")
def step_http_content_length_header(context: Context) -> None:
    start = next(
        msg
        for msg in context.asgi_sent_messages
        if msg.get("type") == "http.response.start"
    )
    headers: list[tuple[bytes, bytes]] = start.get("headers", [])
    header_names = [k.lower() for k, _v in headers]
    assert b"content-length" in header_names, (
        f"Expected content-length header, but found: {headers}"
    )


@then('the HTTP response should include header "{name}" with value "{value}"')
def step_http_header_with_value(context: Context, name: str, value: str) -> None:
    start = next(
        msg
        for msg in context.asgi_sent_messages
        if msg.get("type") == "http.response.start"
    )
    headers: list[tuple[bytes, bytes]] = start.get("headers", [])
    name_bytes = name.lower().encode("utf-8")
    matched_values = [v.decode("utf-8") for k, v in headers if k.lower() == name_bytes]
    assert matched_values, f"Expected {name} header, but none found in: {headers}"
    assert value in matched_values, (
        f"Expected {name} header value {value!r}, got {matched_values!r}"
    )


# --- Unrecognised lifespan message type ---


@when("I run ASGI lifespan with an unrecognised message type")
def step_run_lifespan_with_unrecognised_type(context: Context) -> None:
    sent_messages: list[SendMessage] = []
    incoming = deque(
        [
            {"type": "lifespan.startup"},
            {"type": "lifespan.bogus"},
            {"type": "lifespan.shutdown"},
        ]
    )

    async def receive() -> dict[str, Any]:
        if not incoming:
            raise AssertionError(
                "ASGI app issued an unexpected extra receive() call; "
                "all queued lifespan messages have already been consumed"
            )
        return incoming.popleft()

    async def send(message: SendMessage) -> None:
        sent_messages.append(message)

    scope = {"type": "lifespan"}
    logger_name = "cleveragents.a2a.asgi"
    with _capture_log_records(logger_name, logging.WARNING) as records:
        asyncio.run(context.asgi_app(scope, receive, send))
    context.asgi_sent_messages = sent_messages
    context.asgi_log_records = records


@then("a warning should be logged for the unrecognised lifespan message type")
def step_assert_lifespan_warning_logged(context: Context) -> None:
    records: list[logging.LogRecord] = context.asgi_log_records
    warning_messages = [r.getMessage() for r in records]
    assert any("lifespan.bogus" in msg for msg in warning_messages), (
        f"Expected a warning about 'lifespan.bogus', got: {warning_messages}"
    )


class _capture_log_records:
    """Context manager that captures log records for a named logger."""

    def __init__(self, logger_name: str, level: int) -> None:
        self._logger = logging.getLogger(logger_name)
        self._level = level
        self._handler: logging.Handler | None = None
        self.records: list[logging.LogRecord] = []

    def __enter__(self) -> list[logging.LogRecord]:
        handler = logging.Handler()
        handler.setLevel(self._level)
        handler.emit = self.records.append  # type: ignore[assignment]
        self._handler = handler
        self._logger.addHandler(handler)
        self._logger.setLevel(self._level)
        return self.records

    def __exit__(self, *_args: object) -> None:
        if self._handler is not None:
            self._logger.removeHandler(self._handler)