Spaces:
Running
Running
| from __future__ import annotations | |
| import pytest | |
| from app.core.request_limits import RequestBodyLimitMiddleware | |
| def _scope(headers: list[tuple[bytes, bytes]] | None = None) -> dict: | |
| return { | |
| "type": "http", | |
| "asgi": {"version": "3.0"}, | |
| "http_version": "1.1", | |
| "method": "POST", | |
| "scheme": "http", | |
| "path": "/upload", | |
| "raw_path": b"/upload", | |
| "query_string": b"", | |
| "headers": headers or [], | |
| "client": ("127.0.0.1", 12345), | |
| "server": ("testserver", 80), | |
| } | |
| async def _consume_body(scope, receive, send) -> None: | |
| while True: | |
| message = await receive() | |
| if message["type"] == "http.disconnect": | |
| return | |
| if not message.get("more_body", False): | |
| break | |
| await send({"type": "http.response.start", "status": 200, "headers": []}) | |
| await send({"type": "http.response.body", "body": b"ok", "more_body": False}) | |
| async def test_request_body_limit_rejects_oversized_content_length() -> None: | |
| sent = [] | |
| received = False | |
| async def receive(): | |
| nonlocal received | |
| received = True | |
| return {"type": "http.request", "body": b"", "more_body": False} | |
| async def send(message): | |
| sent.append(message) | |
| middleware = RequestBodyLimitMiddleware(_consume_body, max_bytes=5) | |
| await middleware(_scope([(b"content-length", b"6")]), receive, send) | |
| assert sent[0]["status"] == 413 | |
| assert b"no-store" in dict(sent[0]["headers"])[b"cache-control"] | |
| assert received is False | |
| async def test_request_body_limit_rejects_oversized_chunked_body() -> None: | |
| sent = [] | |
| chunks = iter([ | |
| {"type": "http.request", "body": b"123", "more_body": True}, | |
| {"type": "http.request", "body": b"456", "more_body": True}, | |
| {"type": "http.request", "body": b"", "more_body": False}, | |
| ]) | |
| async def receive(): | |
| return next(chunks) | |
| async def send(message): | |
| sent.append(message) | |
| middleware = RequestBodyLimitMiddleware(_consume_body, max_bytes=5) | |
| await middleware(_scope(), receive, send) | |
| assert sent[0]["status"] == 413 | |
| assert b"Request body too large" in sent[1]["body"] | |
| assert sent[1]["more_body"] is False | |
| async def test_request_body_limit_allows_body_at_limit() -> None: | |
| sent = [] | |
| chunks = iter([ | |
| {"type": "http.request", "body": b"12345", "more_body": False}, | |
| ]) | |
| async def receive(): | |
| return next(chunks) | |
| async def send(message): | |
| sent.append(message) | |
| middleware = RequestBodyLimitMiddleware(_consume_body, max_bytes=5) | |
| await middleware(_scope(), receive, send) | |
| assert sent[0]["status"] == 200 | |
| assert sent[1]["body"] == b"ok" | |