lfnovo--open-notebook/tests/test_max_body_size_middleware.py

"""
Tests for api.middleware.MaxBodySizeMiddleware.

Covers the raw-ASGI unit behavior (Content-Length pre-check, streaming
enforcement for chunked/no-Content-Length requests, CORS interaction) plus
integration through a real FastAPI app and the actual registered app.
"""

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from starlette.middleware.cors import CORSMiddleware

from api.middleware import (
    DEFAULT_MAX_UPLOAD_SIZE_MB,
    MaxBodySizeMiddleware,
    get_max_upload_size_bytes,
)


async def _echo_body_app(scope, receive, send):
    """Minimal ASGI app that reads the full body then returns 200."""
    body = b""
    more_body = True
    while more_body:
        message = await receive()
        body += message.get("body", b"")
        more_body = message.get("more_body", False)
    await send({"type": "http.response.start", "status": 200, "headers": []})
    await send(
        {"type": "http.response.body", "body": f"received {len(body)} bytes".encode()}
    )


def make_scope(headers=None, method="POST", path="/upload"):
    raw_headers = [
        (k.lower().encode(), v.encode()) for k, v in (headers or {}).items()
    ]
    return {
        "type": "http",
        "method": method,
        "path": path,
        "headers": raw_headers,
    }


class FakeReceiveStream:
    """Feeds a list of body chunks as successive http.request ASGI messages."""

    def __init__(self, chunks):
        self._chunks = list(chunks)

    async def __call__(self):
        if not self._chunks:
            return {"type": "http.request", "body": b"", "more_body": False}
        chunk = self._chunks.pop(0)
        return {
            "type": "http.request",
            "body": chunk,
            "more_body": bool(self._chunks),
        }


class CollectingSend:
    def __init__(self):
        self.messages = []

    async def __call__(self, message):
        self.messages.append(message)

    @property
    def status(self):
        for m in self.messages:
            if m["type"] == "http.response.start":
                return m["status"]
        return None

    @property
    def body(self):
        return b"".join(
            m.get("body", b"")
            for m in self.messages
            if m["type"] == "http.response.body"
        )

    def header(self, name: bytes):
        for m in self.messages:
            if m["type"] == "http.response.start":
                for k, v in m["headers"]:
                    if k.lower() == name.lower():
                        return v.decode()
        return None


class TestConfig:
    def test_default_is_100mb(self, monkeypatch):
        monkeypatch.delenv("OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB", raising=False)
        assert DEFAULT_MAX_UPLOAD_SIZE_MB == 100
        assert get_max_upload_size_bytes() == 100 * 1024 * 1024

    def test_reads_env_override(self, monkeypatch):
        monkeypatch.setenv("OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB", "5")
        assert get_max_upload_size_bytes() == 5 * 1024 * 1024

    def test_malformed_env_falls_back_to_default(self, monkeypatch):
        monkeypatch.setenv("OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB", "not-a-number")
        assert get_max_upload_size_bytes() == 100 * 1024 * 1024

    def test_non_positive_env_falls_back_to_default(self, monkeypatch):
        # A zero/negative limit would 413 every request that has a body.
        monkeypatch.setenv("OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB", "0")
        assert get_max_upload_size_bytes() == 100 * 1024 * 1024
        monkeypatch.setenv("OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB", "-5")
        assert get_max_upload_size_bytes() == 100 * 1024 * 1024


class TestMiddlewareAsgiLevel:
    @pytest.mark.asyncio
    async def test_content_length_precheck_rejects_before_app_runs(self):
        app_called = False

        async def inner_app(scope, receive, send):
            nonlocal app_called
            app_called = True
            await _echo_body_app(scope, receive, send)

        middleware = MaxBodySizeMiddleware(inner_app, max_body_size=10)
        scope = make_scope(headers={"content-length": "1000"})
        send = CollectingSend()
        await middleware(scope, FakeReceiveStream([b"x" * 1000]), send)

        assert not app_called, "inner app must not run when Content-Length exceeds the limit"
        assert send.status == 413

    @pytest.mark.asyncio
    async def test_small_body_passes_through(self):
        middleware = MaxBodySizeMiddleware(_echo_body_app, max_body_size=1000)
        scope = make_scope(headers={"content-length": "5"})
        send = CollectingSend()
        await middleware(scope, FakeReceiveStream([b"hello"]), send)

        assert send.status == 200
        assert send.body == b"received 5 bytes"

    @pytest.mark.asyncio
    async def test_streaming_body_without_content_length_is_still_enforced(self):
        """A client can omit Content-Length (chunked transfer-encoding) - the
        middleware must still catch an oversized body as it streams in."""
        middleware = MaxBodySizeMiddleware(_echo_body_app, max_body_size=10)
        scope = make_scope(headers={})  # no content-length
        send = CollectingSend()
        chunks = [b"a" * 5, b"b" * 5, b"c" * 5]  # totals 15 > 10
        await middleware(scope, FakeReceiveStream(chunks), send)

        assert send.status == 413

    @pytest.mark.asyncio
    async def test_malformed_content_length_falls_back_to_streaming_check(self):
        middleware = MaxBodySizeMiddleware(_echo_body_app, max_body_size=10)
        scope = make_scope(headers={"content-length": "not-a-number"})
        send = CollectingSend()
        await middleware(scope, FakeReceiveStream([b"x" * 1000]), send)

        assert send.status == 413  # caught by streaming enforcement instead

    @pytest.mark.asyncio
    async def test_non_http_scope_passes_through_untouched(self):
        calls = []

        async def inner_app(scope, receive, send):
            calls.append(scope["type"])

        middleware = MaxBodySizeMiddleware(inner_app, max_body_size=10)
        await middleware({"type": "lifespan"}, FakeReceiveStream([]), CollectingSend())
        assert calls == ["lifespan"]

    @pytest.mark.asyncio
    async def test_exactly_at_limit_passes(self):
        middleware = MaxBodySizeMiddleware(_echo_body_app, max_body_size=10)
        scope = make_scope(headers={"content-length": "10"})
        send = CollectingSend()
        await middleware(scope, FakeReceiveStream([b"x" * 10]), send)
        assert send.status == 200

    @pytest.mark.asyncio
    async def test_one_byte_over_limit_rejected(self):
        middleware = MaxBodySizeMiddleware(_echo_body_app, max_body_size=10)
        scope = make_scope(headers={"content-length": "11"})
        send = CollectingSend()
        await middleware(scope, FakeReceiveStream([b"x" * 11]), send)
        assert send.status == 413


class TestCorsInteraction:
    @pytest.mark.asyncio
    async def test_cors_headers_present_on_413_when_wrapped_by_cors_middleware(self):
        """MaxBodySizeMiddleware must sit *inside* CORSMiddleware so a
        rejected upload still gets CORS headers (mirrors the ordering used
        in api/main.py)."""
        app = CORSMiddleware(
            MaxBodySizeMiddleware(_echo_body_app, max_body_size=10),
            allow_origins=["http://example.com"],
            allow_credentials=True,
        )
        scope = make_scope(
            headers={"content-length": "1000", "origin": "http://example.com"}
        )
        send = CollectingSend()
        await app(scope, FakeReceiveStream([b"x" * 1000]), send)

        assert send.status == 413
        assert send.header(b"access-control-allow-origin") == "http://example.com"


class TestFastApiIntegration:
    """Drives the middleware through a real FastAPI app + TestClient (real
    HTTP-shaped requests via httpx), not just hand-rolled ASGI messages."""

    @pytest.fixture
    def small_limit_app(self):
        app = FastAPI()

        @app.post("/echo")
        async def echo():
            return {"ok": True}

        app.add_middleware(MaxBodySizeMiddleware, max_body_size=100)
        return app

    def test_small_request_passes(self, small_limit_app):
        client = TestClient(small_limit_app)
        response = client.post("/echo", content=b"x" * 50)
        assert response.status_code == 200

    def test_oversized_request_rejected_with_413(self, small_limit_app):
        client = TestClient(small_limit_app)
        response = client.post("/echo", content=b"x" * 1000)
        assert response.status_code == 413
        assert "exceeds the maximum" in response.json()["detail"]


class TestRealAppWiring:
    def test_real_app_registers_middleware_with_configured_size(self):
        from api.main import MAX_UPLOAD_SIZE_BYTES, app

        matches = [m for m in app.user_middleware if m.cls is MaxBodySizeMiddleware]
        assert len(matches) == 1
        assert matches[0].kwargs["max_body_size"] == MAX_UPLOAD_SIZE_BYTES

    def test_real_app_wraps_it_inside_cors(self):
        """CORS must be outermost so it can attach headers to a 413 from
        MaxBodySizeMiddleware - i.e. CORSMiddleware must be added *after* it."""
        from api.main import app

        # getattr: Middleware.cls is typed as a callable protocol without
        # __name__, but every registered middleware here is a plain class.
        classes = [getattr(m.cls, "__name__", "") for m in app.user_middleware]
        assert classes.index("CORSMiddleware") < classes.index("MaxBodySizeMiddleware")