async def test_asgi_send_http() -> None:
server = MockWebsocket()
await server.asgi_send({
"type": "websocket.http.response.start",
"headers": [(b"X-Header", b"Value")],
"status": 200,
})
# Server must not send a response till the receipt of the first
# body chunk.
assert server.sent_events == []
await server.asgi_send({
"type": "websocket.http.response.body",
"body": b"a",
"more_body": True
})
await server.asgi_send({
"type": "websocket.http.response.body",
"more_body": False
})
server.sent_events[0].headers = list(
server.sent_events[0].headers) # To allow comparison
assert server.sent_events == [
RejectConnection(
status_code=200,
headers=[(b"x-header", b"Value"), (b"server", b"hypercorn")],
has_body=True,
),
RejectData(data=b"a", body_finished=False),
RejectData(data=b"", body_finished=True),
]
def test_handshake_rejection_with_body():
events = _make_handshake_rejection(400, b"Hello")
assert events == [
RejectConnection(
headers=[(b"content-length", b"5")], has_body=True, status_code=400
),
RejectData(body_finished=False, data=b"Hello"),
RejectData(body_finished=True, data=b""),
]
async def _asgi_send_rejection(self, message: dict) -> None:
body_suppressed = suppress_body("GET", self.response["status"])
if self.state == ASGIWebsocketState.HANDSHAKE:
headers = chain(
[
(bytes(key).strip(), bytes(value).strip())
for key, value in self.response["headers"]
],
self.response_headers(),
)
await self.asend(
RejectConnection(
status_code=int(self.response["status"]),
headers=headers,
has_body=not body_suppressed,
)
)
self.state = ASGIWebsocketState.RESPONSE
if not body_suppressed:
await self.asend(
RejectData(
data=bytes(message.get("body", b"")),
body_finished=not message.get("more_body", False),
)
)
if not message.get("more_body", False):
await self.asgi_put({"type": "websocket.disconnect"})
self.state = ASGIWebsocketState.HTTPCLOSED
def test_handshake_rejection():
events = _make_handshake_rejection(400)
assert events == [
RejectConnection(
headers=[(b"connection", b"close")], has_body=True, status_code=400
),
RejectData(body_finished=True, data=b""),
]
def reject_request(self,
status_code: int = 400,
reason: str = None) -> None:
if not isinstance(status_code, int):
raise TypeError('status_code must be an integer')
if reason is not None and not isinstance(reason, str):
raise TypeError('reason must be a string')
if not reason:
self._client.sendall(
self._ws.send(RejectConnection(status_code=status_code)))
else:
data = bytearray(
self._ws.send(
RejectConnection(has_body=True,
headers=[(b'Content-type', b'text/txt')
])))
data.extend(self._ws.send(RejectData(reason.encode())))
self._client.sendall(bytes(data))
def _make_handshake_rejection(
status_code: int, body: Optional[bytes] = None
) -> List[Event]:
client = h11.Connection(h11.CLIENT)
server = WSConnection(SERVER)
nonce = generate_nonce()
server.receive_data(
client.send(
h11.Request(
method="GET",
target="/",
headers=[
(b"Host", b"localhost"),
(b"Connection", b"Keep-Alive, Upgrade"),
(b"Upgrade", b"WebSocket"),
(b"Sec-WebSocket-Version", b"13"),
(b"Sec-WebSocket-Key", nonce),
],
)
)
)
if body is not None:
client.receive_data(
server.send(
RejectConnection(
headers=[(b"content-length", b"%d" % len(body))],
status_code=status_code,
has_body=True,
)
)
)
client.receive_data(server.send(RejectData(data=body)))
else:
client.receive_data(server.send(RejectConnection(status_code=status_code)))
events = []
while True:
event = client.next_event()
events.append(event)
if isinstance(event, h11.EndOfMessage):
return events