async def test_serve(nursery):
task = trio.hazmat.current_task()
server = await nursery.start(serve_websocket, echo_request_handler, HOST, 0,
None)
port = server.port
assert server.port != 0
# The server nursery begins with one task (server.listen).
assert len(nursery.child_tasks) == 1
no_clients_nursery_count = len(task.child_nurseries)
async with open_websocket(HOST, port, RESOURCE, use_ssl=False) as conn:
# The server nursery has the same number of tasks, but there is now
# one additional nested nursery.
assert len(nursery.child_tasks) == 1
assert len(task.child_nurseries) == no_clients_nursery_count + 1
async def test_handshake_exception_before_accept():
''' In #107, a request handler that throws an exception before finishing the
handshake causes the task to hang. The proper behavior is to raise an
exception to the nursery as soon as possible. '''
async def handler(request):
raise Exception()
with pytest.raises(Exception):
with trio.open_nursery() as nursery:
server = await nursery.start(serve_websocket, handler, HOST, 0,
None)
async with open_websocket(HOST, server.port, RESOURCE,
use_ssl=False) as client_ws:
pass
async def test_wrap_server_stream(nursery):
async def handler(stream):
request = await wrap_server_stream(nursery, stream)
server_ws = await request.accept()
async with server_ws:
assert not server_ws.closed
msg = await server_ws.get_message()
assert msg == 'Hello from client!'
assert server_ws.closed
serve_fn = partial(trio.serve_tcp, handler, 0, host=HOST)
listeners = await nursery.start(serve_fn)
port = listeners[0].socket.getsockname()[1]
async with open_websocket(HOST, port, RESOURCE, use_ssl=False) as client:
await client.send_message('Hello from client!')
async def test_serve_handler_nursery(nursery):
task = trio.hazmat.current_task()
async with trio.open_nursery() as handler_nursery:
serve_with_nursery = partial(serve_websocket, echo_request_handler,
HOST, 0, None, handler_nursery=handler_nursery)
server = await nursery.start(serve_with_nursery)
port = server.port
# The server nursery begins with one task (server.listen).
assert len(nursery.child_tasks) == 1
no_clients_nursery_count = len(task.child_nurseries)
async with open_websocket(HOST, port, RESOURCE, use_ssl=False) as conn:
# The handler nursery should have one task in it
# (conn._reader_task).
assert len(handler_nursery.child_tasks) == 1
async def test_server_does_not_close_handshake(nursery):
async def handler(stream):
request = await wrap_server_stream(nursery, stream)
server_ws = await request.accept()
async with server_ws:
await stream.aclose()
with pytest.raises(ConnectionClosed):
await server_ws.send_message('Hello from client!')
serve_fn = partial(trio.serve_tcp, handler, 0, host=HOST)
listeners = await nursery.start(serve_fn)
port = listeners[0].socket.getsockname()[1]
async with open_websocket(HOST, port, RESOURCE, use_ssl=False) as client:
with pytest.raises(ConnectionClosed):
await client.get_message()
async def test_serve_ssl(nursery):
server_context = trio.ssl.create_default_context(
trio.ssl.Purpose.CLIENT_AUTH)
client_context = trio.ssl.create_default_context()
ca = trustme.CA()
ca.configure_trust(client_context)
cert = ca.issue_server_cert(HOST)
cert.configure_cert(server_context)
server = await nursery.start(serve_websocket, echo_request_handler, HOST,
0, server_context)
port = server.port
async with open_websocket(HOST, port, RESOURCE,
use_ssl=client_context) as conn:
assert not conn.closed
async def test_handshake_subprotocol(nursery):
async def handler(request):
assert request.proposed_subprotocols == ('chat', 'file')
assert request.subprotocol is None
request.subprotocol = 'chat'
assert request.subprotocol == 'chat'
server_ws = await request.accept()
assert server_ws.subprotocol == 'chat'
server = await nursery.start(serve_websocket, handler, HOST, 0, None)
async with open_websocket(HOST,
server.port,
RESOURCE,
use_ssl=False,
subprotocols=('chat', 'file')) as client_ws:
assert client_ws.subprotocol == 'chat'
async def test_server_handler_exit(nursery, autojump_clock):
async def handler(request):
server_ws = await request.accept()
await trio.sleep(1)
server = await nursery.start(
partial(serve_websocket, handler, HOST, 0, ssl_context=None))
# connection should close when server handler exists
with trio.fail_after(2):
async with open_websocket(
HOST, server.port, '/', use_ssl=False) as connection:
with pytest.raises(ConnectionClosed) as exc_info:
await connection.get_message()
exc = exc_info.value
assert exc.reason.name == 'NORMAL_CLOSURE'
async def test_reject_handshake(nursery):
async def handler(request):
body = b'My body'
await request.reject(400, body=body)
server = await nursery.start(serve_websocket, handler, HOST, 0, None)
with pytest.raises(ConnectionRejected) as exc_info:
async with open_websocket(
HOST,
server.port,
RESOURCE,
use_ssl=False,
) as client_ws:
pass
exc = exc_info.value
assert exc.body == b'My body'
async def test_handshake_client_headers(nursery):
async def handler(request):
headers = dict(request.headers)
assert b'x-test-header' in headers
assert headers[b'x-test-header'] == b'My test header'
server_ws = await request.accept()
await server_ws.send_message('test')
server = await nursery.start(serve_websocket, handler, HOST, 0, None)
headers = [(b'X-Test-Header', b'My test header')]
async with open_websocket(HOST,
server.port,
RESOURCE,
use_ssl=False,
extra_headers=headers) as client_ws:
await client_ws.get_message()
async def test_client_open_timeout(nursery, autojump_clock):
'''
The client times out waiting for the server to complete the opening
handshake.
'''
async def handler(request):
await trio.sleep(FORCE_TIMEOUT)
server_ws = await request.accept()
pytest.fail('Should not reach this line.')
server = await nursery.start(
partial(serve_websocket, handler, HOST, 0, ssl_context=None))
with pytest.raises(trio.TooSlowError):
async with open_websocket(HOST, server.port, '/', use_ssl=False,
connect_timeout=TIMEOUT) as client_ws:
pass
async def main() -> None:
"""Print all received messages to console."""
logger.info('Opening connection...')
while True:
try:
async with trio_websocket.open_websocket(
'127.0.0.1', 8000, '/', use_ssl=False
) as connection:
logger.info('Opened connection!')
while True:
logger.info('Sending ping message...')
await connection.send_message(b'ping!')
message = await connection.get_message()
logger.info('Received message: %s', message)
except trio_websocket.HandshakeError:
await trio.sleep(1)
except trio_websocket.ConnectionClosed as exc:
logger.info(
'Connection closed! Attempting to reconnect...', exc_info=exc
)
await trio.sleep(1)
async def test_reject_handshake_invalid_info_status(nursery):
'''
An informational status code that is not 101 should cause the client to
reject the handshake. Since it is an informational response, there will not
be a response body, so this test exercises a different code path.
'''
async def handler(stream):
await stream.send_all(b'HTTP/1.1 100 CONTINUE\r\n\r\n')
await stream.receive_some(max_bytes=1024)
serve_fn = partial(trio.serve_tcp, handler, 0, host=HOST)
listeners = await nursery.start(serve_fn)
port = listeners[0].socket.getsockname()[1]
with pytest.raises(ConnectionRejected) as exc_info:
async with open_websocket(HOST, port, RESOURCE, use_ssl=False,
) as client_ws:
pass
exc = exc_info.value
assert exc.status_code == 100
assert repr(exc) == 'ConnectionRejected<status_code=100>'
assert exc.body is None
async def test_no_messages_after_local_close(nursery):
'''
If the local endpoint initiates closing, then pending messages are discarded
and any attempt to read a message will raise ConnectionClosed.
'''
client_closed = trio.Event()
async def handler(request):
# The server sends some messages and then closes.
server = await request.accept()
async with server:
await server.send_message('1')
await server.send_message('2')
await client_closed.wait()
server = await nursery.start(
partial(serve_websocket, handler, HOST, 0, ssl_context=None))
async with open_websocket(HOST, server.port, '/', use_ssl=False) as client:
pass
with pytest.raises(ConnectionClosed):
await client.get_message()
client_closed.set()
async def echo_conn(echo_server):
''' Return a client connection instance that is connected to an echo
server. '''
async with open_websocket(HOST, echo_server.port, RESOURCE,
use_ssl=False) as conn:
await yield_(conn)
async def test_client_open(echo_server):
async with open_websocket(HOST, echo_server.port, RESOURCE, use_ssl=False) \
as conn:
assert not conn.closed
async def client(server):
async with open_websocket(HOST, server.port, '/', use_ssl=False) as ws:
yield ws
async def test_client_open(echo_server):
async with open_websocket(HOST, echo_server.port, RESOURCE, use_ssl=False) \
as conn:
assert not conn.closed
assert conn.is_client
assert str(conn).startswith('client-')
