def test_postgresql_backend(
mock_ws_connect_event,
mock_ws_send_event,
mock_ws_disconnect_event,
mock_websocket_app,
):
with testing.postgresql.Postgresql() as postgresql:
create_engine(postgresql.url())
dsn = postgresql.url()
handler = Mangum(mock_websocket_app, dsn=dsn)
response = handler(mock_ws_connect_event, {})
assert response == {"statusCode": 200}
handler = Mangum(mock_websocket_app, dsn=dsn)
with mock.patch(
"mangum.websocket.WebSocket.post_to_connection") as send:
send.return_value = None
response = handler(mock_ws_send_event, {})
assert response == {"statusCode": 200}
handler = Mangum(mock_websocket_app, dsn=dsn)
response = handler(mock_ws_disconnect_event, {})
assert response == {"statusCode": 200}
handler = Mangum(mock_websocket_app, dsn=dsn)
response = handler(mock_ws_connect_event, {})
assert response == {"statusCode": 200}
def test_http_cycle_state(mock_http_event) -> None:
async def app(scope, receive, send):
assert scope["type"] == "http"
await send({"type": "http.response.body", "body": b"Hello, world!"})
handler = Mangum(app, enable_lifespan=False)
with pytest.raises(RuntimeError):
handler(mock_http_event, {})
async def app(scope, receive, send):
assert scope["type"] == "http"
await send({
"type": "http.response.start",
"status": 200,
"headers": []
})
await send({
"type": "http.response.start",
"status": 200,
"headers": []
})
handler = Mangum(app, enable_lifespan=False)
with pytest.raises(RuntimeError):
handler(mock_http_event, {})
def test_http_api_gateway_base_path(mock_http_event) -> None:
async def app(scope, receive, send):
assert scope["type"] == "http"
assert scope["path"] == urllib.parse.unquote(mock_http_event["path"])
await send({"type": "http.response.start", "status": 200})
await send({"type": "http.response.body", "body": b"Hello world!"})
handler = Mangum(app, lifespan="off", api_gateway_base_path=None)
response = handler(mock_http_event, {})
assert response == {
"body": "Hello world!",
"headers": {},
"isBase64Encoded": False,
"statusCode": 200,
}
async def app(scope, receive, send):
assert scope["type"] == "http"
assert scope["path"] == urllib.parse.unquote(
mock_http_event["path"][len(f"/{api_gateway_base_path}"):])
await send({"type": "http.response.start", "status": 200})
await send({"type": "http.response.body", "body": b"Hello world!"})
api_gateway_base_path = "test"
handler = Mangum(app,
lifespan="off",
api_gateway_base_path=api_gateway_base_path)
response = handler(mock_http_event, {})
assert response == {
"body": "Hello world!",
"headers": {},
"isBase64Encoded": False,
"statusCode": 200,
}
def test_asgi_cycle_state(mock_data) -> None:
def app(scope):
async def asgi(receive, send):
await send({
"type": "http.response.body",
"body": b"Hello, world!"
})
return asgi
mock_event = mock_data.get_aws_event()
with pytest.raises(RuntimeError):
Mangum(app)(mock_event, {})
def app(scope):
async def asgi(receive, send):
await send({
"type": "http.response.start",
"status": 200,
"headers": []
})
await send({
"type": "http.response.start",
"status": 200,
"headers": []
})
return asgi
mock_event = mock_data.get_aws_event()
with pytest.raises(RuntimeError):
Mangum(app)(mock_event, {})
def test_s3_backend(
tmp_path,
mock_ws_connect_event,
mock_ws_send_event,
mock_ws_disconnect_event,
mock_websocket_app,
dsn,
) -> None:
# bucket = "mangum-bucket-12345"
# conn = boto3.resource("s3")
# conn.create_bucket(Bucket=bucket)
handler = Mangum(mock_websocket_app, dsn=dsn)
response = handler(mock_ws_connect_event, {})
assert response == {"statusCode": 200}
handler = Mangum(mock_websocket_app, dsn=dsn)
with mock.patch("mangum.websocket.WebSocket.post_to_connection") as send:
send.return_value = None
response = handler(mock_ws_send_event, {})
assert response == {"statusCode": 200}
handler = Mangum(mock_websocket_app, dsn=dsn)
response = handler(mock_ws_disconnect_event, {})
assert response == {"statusCode": 200}
def test_http_cycle_state(mock_http_event) -> None:
async def app(scope, receive, send):
assert scope["type"] == "http"
await send({"type": "http.response.body", "body": b"Hello, world!"})
handler = Mangum(app, enable_lifespan=False)
response = handler(mock_http_event, {})
assert response == {
"body": "Internal Server Error",
"headers": {
"content-type": "text/plain; charset=utf-8"
},
"isBase64Encoded": False,
"statusCode": 500,
}
async def app(scope, receive, send):
assert scope["type"] == "http"
await send({"type": "http.response.start", "status": 200})
await send({"type": "http.response.start", "status": 200})
handler = Mangum(app, enable_lifespan=False)
response = handler(mock_http_event, {})
assert response == {
"body": "Internal Server Error",
"headers": {
"content-type": "text/plain; charset=utf-8"
},
"isBase64Encoded": False,
"statusCode": 500,
}
def test_websocket_unexpected_message_error(tmp_path, mock_ws_connect_event,
mock_ws_send_event) -> None:
async def app(scope, receive, send):
await send({"type": "websocket.oops", "subprotocol": None})
dsn = f"sqlite://{tmp_path}/mangum.sqlite3"
handler = Mangum(app, dsn=dsn)
handler(mock_ws_connect_event, {})
handler = Mangum(app, dsn=dsn)
response = handler(mock_ws_send_event, {})
assert response == {"statusCode": 500}
def test_websocket_exception(tmp_path, mock_ws_connect_event,
mock_ws_send_event) -> None:
async def app(scope, receive, send):
raise Exception()
dsn = f"sqlite://{tmp_path}/mangum.sqlite3"
handler = Mangum(app, dsn=dsn)
handler(mock_ws_connect_event, {})
handler = Mangum(app, dsn=dsn)
response = handler(mock_ws_send_event, {})
assert response == {"statusCode": 500}
def test_starlette_websocket(mock_ws_connect_event, mock_ws_send_event,
dynamodb) -> None:
table_name = "test-table"
dynamodb.create_table(
TableName=table_name,
KeySchema=[{
"AttributeName": "connectionId",
"KeyType": "HASH"
}],
AttributeDefinitions=[{
"AttributeName": "connectionId",
"AttributeType": "S"
}],
ProvisionedThroughput={
"ReadCapacityUnits": 5,
"WriteCapacityUnits": 5
},
)
app = Starlette()
@app.websocket_route("/ws")
async def websocket_endpoint(websocket):
await websocket.accept()
await websocket.send_json({"url": str(websocket.url)})
await websocket.close()
handler = Mangum(app, enable_lifespan=False)
response = handler(mock_ws_connect_event, {})
assert response == {
"body": "OK",
"headers": {
"content-type": "text/plain; charset=utf-8"
},
"isBase64Encoded": False,
"statusCode": 200,
}
handler = Mangum(app, enable_lifespan=False)
with mock.patch("mangum.protocols.websockets.ASGIWebSocketCycle.send_data"
) as send_data:
send_data.return_value = None
response = handler(mock_ws_send_event, {})
assert response == {
"body": "OK",
"headers": {
"content-type": "text/plain; charset=utf-8"
},
"isBase64Encoded": False,
"statusCode": 200,
}
def test_http_text_mime_types(mock_http_event) -> None:
async def app(scope, receive, send):
assert scope["type"] == "http"
await send({
"type":
"http.response.start",
"status":
200,
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
})
await send({"type": "http.response.body", "body": b"Hello, world!"})
handler = Mangum(app,
enable_lifespan=False,
text_mime_types=["application/vnd.apple.pkpass"])
response = handler(mock_http_event, {})
assert response == {
"statusCode": 200,
"isBase64Encoded": False,
"headers": {
"content-type": "text/plain; charset=utf-8"
},
"body": "Hello, world!",
}
def test_http_empty_header(mock_http_event) -> None:
async def app(scope, receive, send):
assert scope["type"] == "http"
await send({
"type":
"http.response.start",
"status":
200,
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
})
await send({"type": "http.response.body", "body": b"Hello, world!"})
handler = Mangum(app, lifespan="off")
mock_http_event["headers"] = None
response = handler(mock_http_event, {})
assert response == {
"statusCode": 200,
"isBase64Encoded": False,
"headers": {
"content-type": "text/plain; charset=utf-8"
},
"body": "Hello, world!",
}
def test_http_binary_gzip_response(mock_http_event) -> None:
body = json.dumps({"abc": "defg"})
async def app(scope, receive, send):
assert scope["type"] == "http"
await send({
"type": "http.response.start",
"status": 200,
"headers": [[b"content-type", b"application/json"]],
})
await send({"type": "http.response.body", "body": body.encode()})
handler = Mangum(GZipMiddleware(app, minimum_size=1), lifespan="off")
response = handler(mock_http_event, {})
assert response["isBase64Encoded"]
assert response["headers"] == {
"content-encoding": "gzip",
"content-type": "application/json",
"content-length": "35",
"vary": "Accept-Encoding",
}
assert response["body"] == base64.b64encode(gzip.compress(
body.encode())).decode()
def test_aws_binary_response_with_body(mock_data) -> None:
def app(scope):
async def asgi(receive, send):
message = await receive()
body = message["body"]
response = PlainTextResponse(body)
await response(receive, send)
return asgi
mock_event = mock_data.get_aws_event()
body = b"123"
body_encoded = base64.b64encode(body)
mock_event["body"] = body_encoded
mock_event["isBase64Encoded"] = True
handler = Mangum(app)
response = handler(mock_event, {})
assert response == {
"statusCode": 200,
"isBase64Encoded": True,
"headers": {
"content-length": "3",
"content-type": "text/plain; charset=utf-8"
},
"body": body_encoded.decode(),
}
def test_http_binary_request_and_response(mock_http_event) -> None:
async def app(scope, receive, send):
assert scope["type"] == "http"
body = []
message = await receive()
if "body" in message:
body.append(message["body"])
if not message.get("more_body", False):
body = b"".join(body)
await send({
"type":
"http.response.start",
"status":
200,
"headers": [[b"content-type", b"application/octet-stream"]],
})
await send({"type": "http.response.body", "body": b"abc"})
mock_http_event["isBase64Encoded"] = True
handler = Mangum(app, lifespan="off")
response = handler(mock_http_event, {})
assert response == {
"statusCode": 200,
"isBase64Encoded": True,
"headers": {
"content-type": "application/octet-stream"
},
"body": base64.b64encode(b"abc").decode(),
}
def test_http_response_with_body(mock_http_event) -> None:
async def app(scope, receive, send):
assert scope["type"] == "http"
body = [b"4", b"5", b"6"]
while True:
message = await receive()
if "body" in message:
body.append(message["body"])
if not message.get("more_body", False):
body = b"".join(body)
await send({
"type":
"http.response.start",
"status":
200,
"headers":
[[b"content-type", b"text/plain; charset=utf-8"]],
})
await send({"type": "http.response.body", "body": body})
return
handler = Mangum(app, lifespan="off")
response = handler(mock_http_event, {})
assert response == {
"statusCode": 200,
"isBase64Encoded": False,
"headers": {
"content-type": "text/plain; charset=utf-8"
},
"body": "456123",
}
def test_lifespan_supported_with_error(mock_http_event) -> None:
async def app(scope, receive, send):
await receive()
await send({
"type":
"http.response.start",
"status":
200,
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
})
await send({"type": "http.response.body", "body": b"Hello, world!"})
handler = Mangum(app)
assert handler.lifespan.is_supported
assert handler.lifespan.has_error
response = handler(mock_http_event, {})
assert response == {
"statusCode": 200,
"isBase64Encoded": False,
"headers": {
"content-type": "text/plain; charset=utf-8"
},
"body": "Hello, world!",
}
def test_http_exception_handler(mock_http_event) -> None:
path = mock_http_event["path"]
app = Starlette()
@app.exception_handler(Exception)
async def all_exceptions(request, exc):
return PlainTextResponse(content="Error!", status_code=500)
@app.route(path)
def homepage(request):
raise Exception()
return PlainTextResponse("Hello, world!")
handler = Mangum(app)
response = handler(mock_http_event, {})
assert response == {
"body": "Error!",
"headers": {
"content-length": "6",
"content-type": "text/plain; charset=utf-8"
},
"isBase64Encoded": False,
"statusCode": 500,
}
def test_lifespan_startup_error(mock_http_event) -> None:
async def app(scope, receive, send):
if scope["type"] == "lifespan":
while True:
message = await receive()
if message["type"] == "lifespan.startup":
raise Exception("error")
else:
await send({
"type":
"http.response.start",
"status":
200,
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
})
await send({
"type": "http.response.body",
"body": b"Hello, world!"
})
handler = Mangum(app)
assert handler.lifespan.is_supported
assert handler.lifespan.has_error
response = handler(mock_http_event, {})
assert response == {
"statusCode": 200,
"isBase64Encoded": False,
"headers": {
"content-type": "text/plain; charset=utf-8"
},
"body": "Hello, world!",
}
def test_http_response_headers(mock_http_event) -> None:
async def app(scope, receive, send):
assert scope["type"] == "http"
await send({
"type":
"http.response.start",
"status":
200,
"headers": [[b"x-header-1", b"123"], [b"x-header-2", b"456"]],
})
await send({"type": "http.response.body", "body": b"Hello, world!"})
handler = Mangum(app, lifespan="off")
mock_http_event["headers"] = None
response = handler(mock_http_event, {})
assert response == {
"statusCode": 200,
"isBase64Encoded": False,
"headers": {
"x-header-1": "123",
"x-header-2": "456"
},
"body": "Hello, world!",
}
def test_http_binary_response__with_body(mock_http_event) -> None:
async def app(scope, receive, send):
assert scope["type"] == "http"
body = []
message = await receive()
if "body" in message:
body.append(message["body"])
if not message.get("more_body", False):
body = b"".join(body)
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
}
)
await send({"type": "http.response.body", "body": body})
mock_http_event["isBase64Encoded"] = True
handler = Mangum(app, enable_lifespan=False)
response = handler(mock_http_event, {})
assert response == {
"statusCode": 200,
"isBase64Encoded": True,
"headers": {"content-type": "text/plain; charset=utf-8"},
"body": base64.b64encode(b"123").decode(),
}
def handler(event, context):
log_vars_and_event(event, context)
asgi_handler = Mangum(app)
response = asgi_handler(
event, context) # Call the instance with the event arguments
return response
def test_lifespan(mock_http_event, lifespan) -> None:
"""
Test each lifespan option using an application that supports lifespan messages.
* "auto" (default):
Application support for lifespan will be inferred.
Any error that occurs during startup will be logged and the ASGI application
cycle will continue unless a `lifespan.startup.failed` event is sent.
* "on":
Application support for lifespan is explicit.
Any error that occurs during startup will be raised and a 500 response will
be returned.
* "off":
Application support for lifespan should be ignored.
The application will not enter the lifespan cycle context.
"""
startup_complete = False
shutdown_complete = False
async def app(scope, receive, send):
nonlocal startup_complete, shutdown_complete
if scope["type"] == "lifespan":
while True:
message = await receive()
if message["type"] == "lifespan.startup":
await send({"type": "lifespan.startup.complete"})
startup_complete = True
elif message["type"] == "lifespan.shutdown":
await send({"type": "lifespan.shutdown.complete"})
shutdown_complete = True
return
if scope["type"] == "http":
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
}
)
await send({"type": "http.response.body", "body": b"Hello, world!"})
handler = Mangum(app, lifespan=lifespan)
response = handler(mock_http_event, {})
expected = lifespan in ("on", "auto")
assert startup_complete == expected
assert shutdown_complete == expected
assert response == {
"statusCode": 200,
"isBase64Encoded": False,
"headers": {"content-type": "text/plain; charset=utf-8"},
"body": "Hello, world!",
}
def test_quart_lifespan(mock_http_event) -> None:
startup_complete = False
shutdown_complete = False
path = mock_http_event["path"]
app = Quart(__name__)
@app.before_serving
async def on_startup():
nonlocal startup_complete
startup_complete = True
@app.after_serving
async def on_shutdown():
nonlocal shutdown_complete
shutdown_complete = True
@app.route(path)
async def hello():
return "hello world!"
assert not startup_complete
assert not shutdown_complete
handler = Mangum(app)
response = handler(mock_http_event, {})
assert startup_complete
assert shutdown_complete
assert response == {
"statusCode": 200,
"isBase64Encoded": False,
"headers": {"content-length": "12", "content-type": "text/html; charset=utf-8"},
"body": "hello world!",
}
def test_lifespan_error(mock_http_event, lifespan, caplog) -> None:
caplog.set_level(logging.ERROR)
async def app(scope, receive, send):
if scope["type"] == "lifespan":
while True:
message = await receive()
if message["type"] == "lifespan.startup":
raise Exception("error")
else:
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
}
)
await send({"type": "http.response.body", "body": b"Hello, world!"})
handler = Mangum(app, lifespan=lifespan)
response = handler(mock_http_event, {})
assert "Exception in 'lifespan' protocol." in caplog.text
assert response == {
"statusCode": 200,
"isBase64Encoded": False,
"headers": {"content-type": "text/plain; charset=utf-8"},
"body": "Hello, world!",
}
def create_handler(app):
"""Create a handler to use with AWS Lambda if mangum available."""
try:
from mangum import Mangum
return Mangum(app)
except ImportError:
return None
def handler(event, context):
asgi_handler = Mangum(app)
response = asgi_handler(event, context)
response["headers"]["access-control-allow-origin"] = "*"
response["headers"]["access-control-allow-credentials"] = "true"
print(response["headers"])
return response
def test_websocket_cycle_state(mock_ws_connect_event, mock_ws_send_event,
dynamodb) -> None:
table_name = "test-table"
dynamodb.create_table(
TableName=table_name,
KeySchema=[{
"AttributeName": "connectionId",
"KeyType": "HASH"
}],
AttributeDefinitions=[{
"AttributeName": "connectionId",
"AttributeType": "S"
}],
ProvisionedThroughput={
"ReadCapacityUnits": 5,
"WriteCapacityUnits": 5
},
)
async def app(scope, receive, send):
await send({"type": "websocket.send", "text": "Hello world!"})
handler = Mangum(app, enable_lifespan=False)
response = handler(mock_ws_connect_event, {})
assert response == {
"body": "OK",
"headers": {
"content-type": "text/plain; charset=utf-8"
},
"isBase64Encoded": False,
"statusCode": 200,
}
handler = Mangum(app, enable_lifespan=False)
with pytest.raises(RuntimeError):
with mock.patch(
"mangum.protocols.websockets.ASGIWebSocketCycle.send_data"
) as send_data:
send_data.return_value = None
handler(mock_ws_send_event, {})
def handler(event, context):
print("<" * 100)
print(json.dumps(event))
fix_missing_fields(event)
asgi_handler = Mangum(app)
response = asgi_handler(
event, context) # Call the instance with the event arguments
print("DONE")
print(">" * 100)
return response
def test_sqlite_3_backend(
tmp_path,
mock_ws_connect_event,
mock_ws_send_event,
mock_ws_disconnect_event,
mock_websocket_app,
) -> None:
dsn = f"sqlite://{tmp_path}/mangum.sqlite3"
handler = Mangum(mock_websocket_app, dsn=dsn)
response = handler(mock_ws_connect_event, {})
assert response == {"statusCode": 200}
handler = Mangum(mock_websocket_app, dsn=dsn)
with mock.patch("mangum.websocket.WebSocket.post_to_connection") as send:
send.return_value = None
response = handler(mock_ws_send_event, {})
assert response == {"statusCode": 200}
handler = Mangum(mock_websocket_app, dsn=dsn)
response = handler(mock_ws_disconnect_event, {})
assert response == {"statusCode": 200}
def test_http_response(mock_http_event) -> None:
async def app(scope, receive, send):
assert scope["type"] == "http"
assert scope == {
"client": ("192.168.100.1", 0),
"headers": [
[
b"accept",
b"text/html,application/xhtml+xml,application/xml;q=0.9,image/"
b"webp,*/*;q=0.8",
],
[b"accept-encoding", b"gzip, deflate, lzma, sdch, br"],
[b"accept-language", b"en-US,en;q=0.8"],
[b"cloudfront-forwarded-proto", b"https"],
[b"cloudfront-is-desktop-viewer", b"true"],
[b"cloudfront-is-mobile-viewer", b"false"],
[b"cloudfront-is-smarttv-viewer", b"false"],
[b"cloudfront-is-tablet-viewer", b"false"],
[b"cloudfront-viewer-country", b"US"],
[b"host", b"test.execute-api.us-west-2.amazonaws.com"],
[b"upgrade-insecure-requests", b"1"],
[b"x-forwarded-for", b"192.168.100.1, 192.168.1.1"],
[b"x-forwarded-port", b"443"],
[b"x-forwarded-proto", b"https"],
],
"http_version": "1.1",
"method": "GET",
"path": "/test/hello",
"query_string": b"name=me",
"raw_path": None,
"root_path": "",
"scheme": "https",
"server": ("test.execute-api.us-west-2.amazonaws.com", 80),
"type": "http",
}
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
}
)
await send({"type": "http.response.body", "body": b"Hello, world!"})
handler = Mangum(app, enable_lifespan=False)
response = handler(mock_http_event, {})
assert response == {
"statusCode": 200,
"isBase64Encoded": False,
"headers": {"content-type": "text/plain; charset=utf-8"},
"body": "Hello, world!",
}