WebSocketConnection now inherits AsyncResource

.gitignore

@@ -1,3 +1,4 @@
+.pytest_cache
 __pycache__
 dist
 examples/fake.*

README.md

@@ -14,7 +14,7 @@ If you want to help develop `trio-websocket`, clone [the
 repository](https://github.com/hyperiongray/trio-websocket) and run this command
 from the repository root:
 
-    pip install --editable .
+    pip install --editable .[dev]
 
 ## Sample client
 
@@ -39,7 +39,11 @@ example client sends a text message and then disconnects.
 
     trio.run(main)
 
-A more detailed example is in `examples/client.py`.
+A more detailed example is in `examples/client.py`. **Note:** if you want to run
+this example client with SSL, you'll need to install the `trustme` module from
+PyPI (installed automtically if you used the `[dev]` extras when installing
+`trio-websocket`) and then generate a self-signed certificate by running
+`example/generate-cert.py`.
 
 ## Sample server
 
@@ -64,7 +68,8 @@ to each incoming message with an identical outgoing message.
 
     trio.run(main)
 
-A longer example is in `examples/server.py`.
+A longer example is in `examples/server.py`. **See the note above about using
+SSL with the example client.**
 
 ## Integration Testing with Autobahn
 

examples/client.py

@@ -44,7 +44,7 @@ def parse_args():
 async def main(args):
     ''' Main entry point, returning False in the case of logged error. '''
     async with trio.open_nursery() as nursery:
-        logging.info('Connecting to WebSocket…')
+        logging.debug('Connecting to WebSocket…')
         ssl_context = ssl.create_default_context()
         if args.ssl:
             try:
@@ -63,30 +63,37 @@ async def main(args):
         except OSError as ose:
             logging.error('Connection attempt failed: %s', ose)
             return False
-        logging.info('Connected!')
-        while True:
+        logging.debug('Connected!')
+        async with connection:
+            await handle_connection(connection)
+        logging.debug('Connection closed')
+
+
+async def handle_connection(connection):
+    ''' Handle the connection. '''
+    while True:
+        try:
+            logger.debug('top of loop')
             await trio.sleep(0.1) # allow time for connection logging
-            try:
-                cmd = await trio.run_sync_in_worker_thread(input, 'cmd> ',
-                    cancellable=True)
-                if cmd.startswith('ping '):
-                    await connection.ping(cmd[5:].encode('utf8'))
-                elif cmd.startswith('send '):
-                    await connection.send_message(cmd[5:])
-                    message = await connection.get_message()
-                    print('response> {}'.format(message))
-                elif cmd.startswith('close'):
-                    try:
-                        reason = cmd[6:]
-                    except IndexError:
-                        reason = None
-                    await connection.close(reason=reason)
-                    break
-                else:
-                    commands()
-            except ConnectionClosed:
-                logging.info('Connection closed')
+            cmd = await trio.run_sync_in_worker_thread(input, 'cmd> ',
+                cancellable=True)
+            if cmd.startswith('ping '):
+                await connection.ping(cmd[5:].encode('utf8'))
+            elif cmd.startswith('send '):
+                await connection.send_message(cmd[5:])
+                message = await connection.get_message()
+                print('response> {}'.format(message))
+            elif cmd.startswith('close'):
+                try:
+                    reason = cmd[6:]
+                except IndexError:
+                    reason = None
+                await connection.aclose(code=1000, reason=reason)
                 break
+            else:
+                commands()
+        except ConnectionClosed:
+            break
 
 
 if __name__ == '__main__':

pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+trio_mode = true

setup.py

@@ -28,9 +28,9 @@
     ],
     keywords='websocket client server trio',
     packages=find_packages(exclude=['docs', 'examples', 'tests']),
-    install_requires=['trio', 'trustme', 'wsproto'],
+    install_requires=['trio', 'wsaccel', 'wsproto'],
     extras_require={
-        'wsaccel': ['wsaccel'],
+        'dev': ['pytest', 'pytest-trio', 'trustme'],
     },
     project_urls={
         'Bug Reports': 'https://github.com/HyperionGray/trio-websocket/issues',

tests/test_connection.py

@@ -0,0 +1,55 @@
+import pytest
+from trio_websocket import ConnectionClosed, WebSocketClient, WebSocketServer
+import trio
+
+
+import logging
+logging.basicConfig(level=logging.DEBUG)
+HOST = 'localhost'
+RESOURCE = 'resource'
+
+
+@pytest.fixture
+async def echo_server(nursery):
+    ''' An echo server reads one message, sends back the same message,
+    then exits. '''
+    async def handler(conn):
+        try:
+            msg = await conn.get_message()
+            await conn.send_message(msg)
+        except ConnectionClosed:
+            pass
+    server = WebSocketServer(handler, HOST, 0, ssl_context=None)
+    await nursery.start(server.listen)
+    yield server
+
+
+@pytest.fixture
+async def echo_conn(echo_server, nursery):
+    ''' Return a client connection instance that is connected to an echo
+    server. '''
+    client = WebSocketClient(HOST, echo_server.port, RESOURCE, use_ssl=False)
+    async with await client.connect(nursery) as conn:
+        yield conn
+
+
+async def test_client_send_and_receive(echo_conn, nursery):
+    async with echo_conn:
+        await echo_conn.send_message('This is a test message.')
+        received_msg = await echo_conn.get_message()
+        assert received_msg == 'This is a test message.'
+
+
+async def test_client_default_close(echo_conn, nursery):
+    async with echo_conn:
+        assert echo_conn.closed is None
+    assert echo_conn.closed.code == 1000
+    assert echo_conn.closed.reason is None
+
+
+async def test_client_nondefault_close(echo_conn, nursery):
+    async with echo_conn:
+        assert echo_conn.closed is None
+        await echo_conn.aclose(code=1001, reason='test reason')
+    assert echo_conn.closed.code == 1001
+    assert echo_conn.closed.reason == 'test reason'