mangum/tests/test_websockets.py at 6f992345d3811f92e51fb07e3df9a3406ea1414a · aminalaee/mangum · GitHub

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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import respx

from mangum import Mangum


@respx.mock(assert_all_mocked=False)
def test_websocket_close(
    sqlite3_dsn, mock_ws_connect_event, mock_ws_send_event
) -> None:
    async def app(scope, receive, send):
        if scope["type"] == "websocket":
            while True:
                message = await receive()
                if message["type"] == "websocket.connect":
                    await send({"type": "websocket.close"})

    handler = Mangum(app, lifespan="off", dsn=sqlite3_dsn)
    response = handler(mock_ws_connect_event, {})
    assert response == {"statusCode": 200}

    response = handler(mock_ws_send_event, {})
    assert response == {"statusCode": 403}


@respx.mock(assert_all_mocked=False)
def test_websocket_disconnect(
    sqlite3_dsn, mock_ws_connect_event, mock_ws_send_event, mock_websocket_app
) -> None:
    handler = Mangum(mock_websocket_app, lifespan="off", dsn=sqlite3_dsn)
    response = handler(mock_ws_connect_event, {})
    assert response == {"statusCode": 200}

    response = handler(mock_ws_send_event, {})
    assert response == {"statusCode": 200}


def test_websocket_exception(
    sqlite3_dsn, mock_ws_connect_event, mock_ws_send_event
) -> None:
    async def app(scope, receive, send):
        raise Exception()

    handler = Mangum(app, dsn=sqlite3_dsn)
    handler(mock_ws_connect_event, {})

    handler = Mangum(app, dsn=sqlite3_dsn)
    response = handler(mock_ws_send_event, {})
    assert response == {"statusCode": 500}


def test_websocket_unexpected_message_error(
    sqlite3_dsn, mock_ws_connect_event, mock_ws_send_event
) -> None:
    async def app(scope, receive, send):
        await send({"type": "websocket.oops", "subprotocol": None})

    handler = Mangum(app, dsn=sqlite3_dsn)
    handler(mock_ws_connect_event, {})

    handler = Mangum(app, dsn=sqlite3_dsn)
    response = handler(mock_ws_send_event, {})
    assert response == {"statusCode": 500}


@respx.mock(assert_all_mocked=False)
def test_websocket_without_body(
    sqlite3_dsn, mock_ws_connect_event, mock_ws_send_event, mock_websocket_app
) -> None:
    handler = Mangum(mock_websocket_app, lifespan="off", dsn=sqlite3_dsn)
    response = handler(mock_ws_connect_event, {})
    assert response == {"statusCode": 200}

    del mock_ws_send_event["body"]
    response = handler(mock_ws_send_event, {})
    assert response == {"statusCode": 200}


@respx.mock(assert_all_mocked=False)
def test_base64_encoded_body_on_request(
    sqlite3_dsn, mock_ws_connect_event, mock_ws_send_event, mock_websocket_app
):
    handler = Mangum(mock_websocket_app, dsn=sqlite3_dsn)
    response = handler(mock_ws_connect_event, {})
    assert response == {"statusCode": 200}

    mock_ws_send_event["body"] = b"bWFuZ3Vt="
    mock_ws_send_event["isBase64Encoded"] = True
    response = handler(mock_ws_send_event, {})
    assert response == {"statusCode": 200}


def test_binary_response(
    sqlite3_dsn, mock_ws_connect_event, mock_ws_send_event, mock_websocket_app
):
    async def app(scope, receive, send):
        if scope["type"] == "websocket":
            while True:
                message = await receive()
                if message["type"] == "websocket.connect":
                    await send({"type": "websocket.accept"})
                elif message["type"] == "websocket.receive":
                    await send({"type": "websocket.send", "body": b"bWFuZ3Vt="})

    handler = Mangum(app, dsn=sqlite3_dsn)
    response = handler(mock_ws_connect_event, {})
    assert response == {"statusCode": 200}

    response = handler(mock_ws_send_event, {})
    assert response == {"statusCode": 500}


@respx.mock(assert_all_mocked=False)
def test_websocket_logging(
    sqlite3_dsn, mock_ws_connect_event, mock_ws_send_event, mock_websocket_app, caplog
) -> None:
    handler = Mangum(mock_websocket_app, lifespan="off", dsn=sqlite3_dsn)
    response = handler(mock_ws_connect_event, {})
    assert response == {"statusCode": 200}

    del mock_ws_send_event["body"]
    response = handler(mock_ws_send_event, {})
    assert response == {"statusCode": 200}

    assert '"WS / HTTP 1.1" 200 0' in caplog.text