use aiohttp.websockets instead of websockets

Author: csko

.travis.yml

@@ -22,7 +22,7 @@ install:
 
 script:
   - py.test --cov=gdax --cov-report=term --cov-append tests
-  
+
 after_success:
   - codecov
 

TODO.md

@@ -1,7 +1,6 @@
 # TODO
 
 - improve test coverage, especially for the order book
-- use aiohttp.websockets instead of websockets
 - client-side rate limiting
 - convert return API values from string to Decimal
 - better enforce API rules

gdax/orderbook.py

@@ -13,7 +13,8 @@
 
 from bintrees import FastRBTree
 import aiofiles
-import websockets
+import aiohttp
+# import websockets
 
 import gdax.trader
 import gdax.utils
@@ -43,11 +44,14 @@ def __init__(self, product_ids='ETH-USD', api_key=None, api_secret=None,
         self._asks = {product_id: FastRBTree() for product_id in product_ids}
         self._bids = {product_id: FastRBTree() for product_id in product_ids}
         self._sequences = {product_id: None for product_id in product_ids}
-        self._ws = None
+        self._ws_session = None
         self._ws_connect = None
+        self._ws = None
 
     async def _init(self):
-        self._ws_connect = websockets.connect('wss://ws-feed.gdax.com')
+        self._ws_session = aiohttp.ClientSession()
+        self._ws_connect = self._ws_session.ws_connect(
+            'wss://ws-feed.gdax.com')
         self._ws = await self._ws_connect.__aenter__()
 
         # subscribe
@@ -89,7 +93,7 @@ async def __aenter__(self):
         return self
 
     async def __aexit__(self, exc_type, exc, traceback):
-        return await self._ws_connect.__aexit__(exc_type, exc, traceback)
+        return await self._ws_session.__aexit__(exc_type, exc, traceback)
 
     async def _open_log_file(self):
         if self.trade_log_file_path is not None:
@@ -101,10 +105,10 @@ async def _close_log_file(self):
             await self._trade_file.__aexit__(None, None, None)
 
     async def _send(self, **kwargs):
-        await self._ws.send(json.dumps(kwargs))
+        await self._ws.send_json(kwargs)
 
     async def _recv(self):
-        json_data = await self._ws.recv()
+        json_data = await self._ws.receive_str()
         if self._trade_file:
             await self._trade_file.write(f'W {json_data}\n')
         return json.loads(json_data)
@@ -132,9 +136,11 @@ async def _subscribe(self):
     async def handle_message(self):
         try:
             message = await self._recv()
-        except websockets.exceptions.ConnectionClosed:
-            await self._ws_connect.__aexit__(None, None, None)
-            self._init()
+        except aiohttp.ServerDisconnectedError as exc:
+            logging.error(
+                f'Error: Exception: f{exc}. Re-initializing websocket.')
+            await self._ws_session.__aexit__(None, None, None)
+            await self._init()
             return
 
         product_id = message['product_id']
@@ -146,10 +152,10 @@ async def handle_message(self):
             # from getProductOrderBook)
             return message
         elif sequence > self._sequences[product_id] + 1:
-            logging.info(
+            logging.error(
                 'Error: messages missing ({} - {}). Re-initializing websocket.'
                 .format(sequence, self._sequences[product_id]))
-            await self._ws_connect.__aexit__(None, None, None)
+            await self._ws_session.__aexit__(None, None, None)
             await self._init()
             return
 
@@ -243,7 +249,6 @@ def match(self, product_id, order):
                 self.set_asks(product_id, price, asks)
 
     def change(self, product_id, order):
-        logging.info((product_id, order))
         if 'new_size' not in order:
             # market order
             # TODO

requirements.txt

@@ -1,6 +1,5 @@
+aiofiles==0.3.1
 aiohttp==2.2.0
 async-timeout==1.2.1
 requests==2.18.1
-websockets==3.3
-aiofiles==0.3.1
 bintrees==2.0.7

setup.py

@@ -12,9 +12,9 @@
       author_email='[email protected]',
       install_requires=[
         'aiohttp==2.2.0',
+        'aiofiles==0.3.1',
         'async-timeout==1.2.1',
         'requests==2.18.1',
-        'websockets==3.3',
         'bintrees==2.0.7',
       ],
       packages=find_packages(),

tests/test_orderbook.py

@@ -138,20 +138,19 @@ def generate_id():
 
 
 @pytest.mark.asyncio
-@patch('websockets.connect', new_callable=AsyncContextManagerMock)
+@patch('aiohttp.ClientSession.ws_connect',
+       new_callable=AsyncContextManagerMock)
 class TestOrderbook(object):
     @patch('gdax.trader.Trader.get_product_order_book')
     async def test_basic_init(self, mock_book, mock_connect):
-        mock_connect.return_value.aenter.send = CoroutineMock()
+        mock_connect.return_value.aenter.send_json = CoroutineMock()
         mock_book.return_value = test_book
 
         product_id = 'ETH-USD'
         product_ids = [product_id]
         async with gdax.orderbook.OrderBook(product_ids) as orderbook:
             msg = {'type': 'subscribe', 'product_ids': product_ids}
-            subscribe_msg = json.dumps(msg)
-            mock_connect.return_value.aenter.send.assert_called_with(
-                subscribe_msg)
+            mock_connect.return_value.aenter.send_json.assert_called_with(msg)
 
             mock_book.assert_called_with(level=3)
 
@@ -195,15 +194,13 @@ async def test_basic_init(self, mock_book, mock_connect):
 
     @patch('gdax.trader.Trader.get_product_order_book')
     async def test_heartbeat(self, mock_book, mock_connect):
-        mock_connect.return_value.aenter.send = CoroutineMock()
+        mock_connect.return_value.aenter.send_json = CoroutineMock()
 
         mock_book.return_value = {'bids': [], 'asks': [], 'sequence': 1}
         product_ids = ['ETH-USD']
         async with gdax.orderbook.OrderBook(product_ids,
                                             use_heartbeat=True) as orderbook:
-            msg1 = {'type': 'subscribe', 'product_ids': product_ids}
-            msg2 = {'type': 'heartbeat', 'on': True}
-            subscribe_msg = json.dumps(msg1)
-            heartbeat_msg = json.dumps(msg2)
+            subscribe_msg = {'type': 'subscribe', 'product_ids': product_ids}
+            heartbeat_msg = {'type': 'heartbeat', 'on': True}
             calls = [call(subscribe_msg), call(heartbeat_msg)]
-            mock_connect.return_value.aenter.send.assert_has_calls(calls)
+            mock_connect.return_value.aenter.send_json.assert_has_calls(calls)