def test_closure(client_sends: bool, code: CloseReason, reason: str) -> None:
client = Connection(CLIENT)
server = Connection(SERVER)
if client_sends:
local = client
remote = server
else:
local = server
remote = client
remote.receive_data(local.send(CloseConnection(code=code, reason=reason)))
event = next(remote.events())
assert isinstance(event, CloseConnection)
assert event.code is code
assert event.reason == reason
assert remote.state is ConnectionState.REMOTE_CLOSING
assert local.state is ConnectionState.LOCAL_CLOSING
local.receive_data(remote.send(event.response()))
event = next(local.events())
assert isinstance(event, CloseConnection)
assert event.code is code
assert event.reason == reason
assert remote.state is ConnectionState.CLOSED # type: ignore[comparison-overlap]
assert local.state is ConnectionState.CLOSED
with pytest.raises(LocalProtocolError):
local.receive_data(b"foobar")
def test_abnormal_closure() -> None:
client = Connection(CLIENT)
client.receive_data(None)
event = next(client.events())
assert isinstance(event, CloseConnection)
assert event.code is CloseReason.ABNORMAL_CLOSURE
assert client.state is ConnectionState.CLOSED
def test_closure(client_sends, code, reason):
client = Connection(CLIENT)
server = Connection(SERVER)
if client_sends:
local = client
remote = server
else:
local = server
remote = client
remote.receive_data(local.send(CloseConnection(code=code, reason=reason)))
event = next(remote.events())
assert isinstance(event, CloseConnection)
assert event.code is code
assert event.reason == reason
assert remote.state is ConnectionState.REMOTE_CLOSING
assert local.state is ConnectionState.LOCAL_CLOSING
local.receive_data(remote.send(event.response()))
event = next(local.events())
assert isinstance(event, CloseConnection)
assert event.code is code
assert event.reason == reason
assert remote.state is ConnectionState.CLOSED
assert local.state is ConnectionState.CLOSED
def test_frame_protocol_gets_fed_garbage() -> None:
client = Connection(CLIENT)
payload = b"x" * 23
frame = b"\x09" + bytearray([len(payload)]) + payload
client.receive_data(frame)
event = next(client.events())
assert isinstance(event, CloseConnection)
assert event.code == CloseReason.PROTOCOL_ERROR
def accept(
self, subprotocol: Optional[str]
) -> Tuple[int, List[Tuple[bytes, bytes]], Connection]:
headers = []
if subprotocol is not None:
if subprotocol not in self.subprotocols:
raise Exception("Invalid Subprotocol")
else:
headers.append((b"sec-websocket-protocol", subprotocol.encode()))
extensions = [PerMessageDeflate()]
accepts = None
if False and self.extensions is not None:
accepts = server_extensions_handshake(self.extensions, extensions)
if accepts:
headers.append((b"sec-websocket-extensions", accepts))
if self.key is not None:
headers.append((b"sec-websocket-accept", generate_accept_token(self.key)))
status_code = 200
if self.http_version == "1.1":
headers.extend([(b"upgrade", b"WebSocket"), (b"connection", b"Upgrade")])
status_code = 101
return status_code, headers, Connection(ConnectionType.SERVER, extensions)
def test_unsolicited_pong() -> None:
client = Connection(CLIENT)
server = Connection(SERVER)
payload = b"x" * 23
server.receive_data(client.send(Pong(payload=payload)))
event = next(server.events())
assert isinstance(event, Pong)
assert event.payload == payload
def test_send_message(client_sends, final):
client = Connection(CLIENT)
server = Connection(SERVER)
if client_sends:
local = client
remote = server
else:
local = server
remote = client
data = b"x" * 23
remote.receive_data(local.send(Message(data=data, message_finished=final)))
event = next(remote.events())
assert isinstance(event, BytesMessage)
assert event.data == data
assert event.message_finished is final
def test_data(split_message: bool) -> None:
client = Connection(CLIENT)
server = Connection(SERVER)
data = "ƒñö®∂😎"
server.receive_data(
client.send(TextMessage(data=data,
message_finished=not split_message)))
event = next(server.events())
assert isinstance(event, TextMessage)
assert event.message_finished is not split_message
def test_ping_pong(client_sends: bool) -> None:
client = Connection(CLIENT)
server = Connection(SERVER)
if client_sends:
local = client
remote = server
else:
local = server
remote = client
payload = b"x" * 23
remote.receive_data(local.send(Ping(payload=payload)))
event = next(remote.events())
assert isinstance(event, Ping)
assert event.payload == payload
local.receive_data(remote.send(event.response()))
event = next(local.events())
assert isinstance(event, Pong)
assert event.payload == payload
def accept(
self,
subprotocol: Optional[str],
additional_headers: Iterable[Tuple[bytes, bytes]],
) -> Tuple[int, List[Tuple[bytes, bytes]], Connection]:
headers = []
if subprotocol is not None:
if subprotocol not in self.subprotocols:
raise Exception("Invalid Subprotocol")
else:
headers.append(
(b"sec-websocket-protocol", subprotocol.encode()))
extensions: List[Extension] = [PerMessageDeflate()]
accepts = None
if self.extensions is not None:
accepts = server_extensions_handshake(self.extensions, extensions)
if accepts:
headers.append((b"sec-websocket-extensions", accepts))
if self.key is not None:
headers.append(
(b"sec-websocket-accept", generate_accept_token(self.key)))
status_code = 200
if self.http_version == "1.1":
headers.extend([(b"upgrade", b"WebSocket"),
(b"connection", b"Upgrade")])
status_code = 101
for name, value in additional_headers:
if b"sec-websocket-protocol" == name or name.startswith(b":"):
raise Exception(f"Invalid additional header, {name.decode()}")
headers.append((name, value))
return status_code, headers, Connection(ConnectionType.SERVER,
extensions)
def test_close_whilst_closing() -> None:
client = Connection(CLIENT)
client.send(CloseConnection(code=CloseReason.NORMAL_CLOSURE))
with pytest.raises(LocalProtocolError):
client.send(CloseConnection(code=CloseReason.NORMAL_CLOSURE))
def test_receive_data_when_closed() -> None:
client = Connection(CLIENT)
client._state = ConnectionState.CLOSED
with pytest.raises(LocalProtocolError):
client.receive_data(b"something")
def test_send_invalid_event() -> None:
client = Connection(CLIENT)
with pytest.raises(LocalProtocolError):
client.send(Request(target="/", host="wsproto"))
