From 37388d8122960f6aec707a9aa30a310bdc67610a Mon Sep 17 00:00:00 2001
From: John Belmonte <john@neggie.net>
Date: Sat, 21 Nov 2020 15:38:54 +0900
Subject: [PATCH] use native async generator support

---
 tests/test_connection.py | 7 ++-----
 trio_websocket/_impl.py  | 5 ++---
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/tests/test_connection.py b/tests/test_connection.py
index 8af07a9..7832355 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -37,7 +37,6 @@
 import pytest
 import trio
 import trustme
-from async_generator import async_generator, yield_
 from trio.testing import memory_stream_pair
 from wsproto.events import CloseConnection
 
@@ -78,24 +77,22 @@
 
 
 @pytest.fixture
-@async_generator
 async def echo_server(nursery):
     ''' A server that reads one message, sends back the same message,
     then closes the connection. '''
     serve_fn = partial(serve_websocket, echo_request_handler, HOST, 0,
         ssl_context=None)
     server = await nursery.start(serve_fn)
-    await yield_(server)
+    yield server
 
 
 @pytest.fixture
-@async_generator
 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)
+        yield conn
 
 
 async def echo_request_handler(request):
diff --git a/trio_websocket/_impl.py b/trio_websocket/_impl.py
index 693c379..440d0e2 100644
--- a/trio_websocket/_impl.py
+++ b/trio_websocket/_impl.py
@@ -9,7 +9,7 @@
 import struct
 import urllib.parse
 
-from async_generator import async_generator, yield_, asynccontextmanager
+from async_generator import asynccontextmanager
 import trio
 import trio.abc
 from wsproto import ConnectionType, WSConnection
@@ -65,7 +65,6 @@ def remove_cancels(exc):
 
 
 @asynccontextmanager
-@async_generator
 async def open_websocket(host, port, resource, *, use_ssl, subprotocols=None,
     extra_headers=None,
     message_queue_size=MESSAGE_QUEUE_SIZE, max_message_size=MAX_MESSAGE_SIZE,
@@ -115,7 +114,7 @@ async def open_websocket(host, port, resource, *, use_ssl, subprotocols=None,
         except OSError as e:
             raise HandshakeError from e
         try:
-            await yield_(connection)
+            yield connection
         finally:
             try:
                 with trio.fail_after(disconnect_timeout):