Spaces:
Running
Running
File size: 2,792 Bytes
03becae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | 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})
@pytest.mark.asyncio
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
@pytest.mark.asyncio
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
@pytest.mark.asyncio
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"
|