async def test_lifespan_default_callbacks_start_are_used_when_nothing_is_provided(
mocked_on_start):
send_data = []
called_once = False
async def receive_start_up():
nonlocal called_once
if not called_once:
called_once = True
return {"type": "lifespan.startup"}
else:
await asyncio.sleep(10)
async def send_mocked(arg):
send_data.append(arg)
server = build_server(handler_identity,
on_start=awaitable_mock(mocked_on_start))
lifecycle_handler = server({"type": "lifespan"})
from asyncio import TimeoutError as AsyncTimeoutError
with pytest.raises(AsyncTimeoutError):
await asyncio.wait_for(asyncio.ensure_future(
lifecycle_handler(receive_start_up, send_mocked)),
timeout=0.3)
mocked_on_start.assert_called_once()
assert send_data == [{"type": "lifespan.startup.complete"}]
async def test_ws_handler_can_have_custom_on_disconnect():
disconnect = Mock()
_ = await run_ws_handler_with_messages(websocket(
disconnect_immediatly, on_disconnect=awaitable_mock(disconnect)),
[{
"type": "websocket.connect"
}],
scope_type="websocket")
disconnect.assert_called_once()
async def test_ws_handler_can_have_custom_on_close():
close = Mock()
_ = await run_ws_handler_with_messages(websocket(
close_immediatly, on_close=awaitable_mock(close)),
[{
"type": "websocket.connect"
}],
scope_type="websocket")
close.assert_called_once()
async def run_ws_handler_with_messages(handler,
messages,
scope_type="websocket"):
scope = {"type": scope_type}
sender = Mock().send
async_sender = awaitable_mock(sender)
receiver = receive_func_from_coll(messages)
initialized_handler = await handler(scope)
await initialized_handler(receiver, async_sender)
return sender.call_args_list
async def test_gen_strategy_will_send_every_message_that_gets_yielded():
async def ws_echo_handler(scope, receiver):
async for msg in receiver:
yield ws_receive(msg)
sender = Mock().send
async_sender = awaitable_mock(sender)
fake_messages = list(map(ws_receive, ["a", b"", "1", b"2"]))
receiver = receive_func_from_coll(fake_messages)
await _ws_async_generator_client(ws_echo_handler, {}, tuple(), receiver,
async_sender)
sender.assert_has_calls(list(map(call, fake_messages)))
async def test_close_will_only_get_once_when_client_yields_a_close():
async def ws_close_actively(scope, receiver):
yield ws_close(2000)
sender = Mock().send
async_sender = awaitable_mock(sender)
fake_messages = list(map(ws_receive, ["a"]))
receiver = receive_func_from_coll(fake_messages)
await _ws_async_generator_client(ws_close_actively, {}, tuple(), receiver,
async_sender)
assert sender.call_args_list == list(
map(call, [{
"type": "websocket.close",
"code": 2000
}]))
async def test_close_gets_automatically_sended_when_client_stops_yielding_messages(
):
async def ws_close_immediately(scope, receiver):
yield ws_send("adf")
sender = Mock().send
async_sender = awaitable_mock(sender)
fake_messages = list(map(ws_receive, ["a"]))
receiver = receive_func_from_coll(fake_messages)
await _ws_async_generator_client(ws_close_immediately, {}, tuple(),
receiver, async_sender)
assert sender.call_args_list == list(
map(call, [ws_send("adf"), {
"type": "websocket.close",
"code": 1000
}]))
async def test_lifespan_default_callbacks_stops_are_used_when_nothing_is_provided(
mocked_on_stop):
send_data = []
async def receive_shutdown():
return {"type": "lifespan.shutdown"}
async def send_mocked(arg):
send_data.append(arg)
server = build_server(handler_identity,
on_stop=awaitable_mock(mocked_on_stop))
lifecycle_handler = server({"type": "lifespan"})
await lifecycle_handler(receive_shutdown, send_mocked)
mocked_on_stop.assert_called_once()
assert send_data == [{"type": "lifespan.shutdown.complete"}]
async def test_ws_handler_can_be_non_generators():
async def ws_just_return_first_message(scope, receiver):
async for msg in receiver:
return ws_receive(msg)
async def ws_return_nothing(scope, receiver):
pass
sender = Mock().send
async_sender = awaitable_mock(sender)
fake_messages = list(map(ws_receive, ["a", b"", "1", b"2"]))
receiver = receive_func_from_coll(fake_messages)
await _ws_async_generator_client(ws_just_return_first_message, {}, tuple(),
receiver, async_sender)
assert sender.mock_calls == [call(fake_messages[0]), call(ws_close())]
sender.reset_mock()
await _ws_async_generator_client(ws_return_nothing, {}, tuple(), receiver,
async_sender)
assert sender.mock_calls == [call(ws_close())]
def receive_func_from_coll(coll):
mock = Mock()
mock.side_effect = list(coll) + [StopAsyncIteration]
return awaitable_mock(mock)