fix websocket connection error handling and add tests

Author: csko

gdax/orderbook.py

@@ -14,7 +14,6 @@
 from bintrees import FastRBTree
 import aiofiles
 import aiohttp
-# import websockets
 
 import gdax.trader
 import gdax.utils
@@ -106,7 +105,7 @@ async def __aexit__(self, exc_type, exc, traceback):
     async def _open_log_file(self):
         if self.trade_log_file_path is not None:
             self._trade_file = await aiofiles.open(self.trade_log_file_path,
-                                                   mode='w').__aenter__()
+                                                   mode='a').__aenter__()
 
     async def _close_log_file(self):
         if self._trade_file is not None:
@@ -147,8 +146,8 @@ async def handle_message(self):
         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()
+            await self.__aexit__(None, None, None)
+            await self.__aenter__()
             return
 
         msg_type = message['type']
@@ -168,8 +167,8 @@ async def handle_message(self):
             logging.error(
                 'Error: messages missing ({} - {}). Re-initializing websocket.'
                 .format(sequence, self._sequences[product_id]))
-            await self._ws_session.__aexit__(None, None, None)
-            await self._init()
+            await self.__aexit__(None, None, None)
+            await self.__aenter__()
             return
 
         if msg_type == 'open':

tests/test_orderbook.py

@@ -4,6 +4,7 @@
 from decimal import Decimal
 
 import asyncio
+import aiohttp
 import pytest
 from asynctest import MagicMock, patch, CoroutineMock, call
 
@@ -457,7 +458,7 @@ async def test_orderbook_advanced(self, mock_book, mock_connect):
             # TODO
 
     @patch('gdax.trader.Trader.get_product_order_book')
-    async def test_error(self, mock_book, mock_connect):
+    async def test_error_message(self, mock_book, mock_connect):
         mock_connect.return_value.aenter.receive_str = CoroutineMock()
         mock_connect.return_value.aenter.send_json = CoroutineMock()
         mock_book.return_value = {'bids': [], 'asks': [], 'sequence': 1}
@@ -482,3 +483,101 @@ async def test_error(self, mock_book, mock_connect):
 
             with pytest.raises(gdax.orderbook.OrderBookError):
                 message = await orderbook.handle_message()
+
+    @patch('gdax.trader.Trader.get_product_order_book')
+    async def test_disconnect(self, mock_book, mock_connect):
+        mock_connect.return_value.aenter.receive_str = CoroutineMock()
+        mock_connect.return_value.aenter.send_json = CoroutineMock()
+        mock_book.return_value = {'bids': [], 'asks': [], 'sequence': 1}
+
+        messages_expected = [
+            json.dumps({
+              "type": "done",
+              "side": "sell",
+              "order_id": "4eef1226-4b38-422c-a5b1-56def7107f9a",
+              "reason": "canceled",
+              "product_id": "ETH-USD",
+              "price": "2601.76000000",
+              "remaining_size": "3.09000000",
+              "sequence": 2,
+              "time": "2017-06-25T11:23:14.775000Z"
+            }),
+            aiohttp.ServerDisconnectedError('error'),
+            json.dumps({
+              "type": "done",
+              "side": "sell",
+              "order_id": "4eef1226-4b38-422c-a5b1-56def7107f9a",
+              "reason": "canceled",
+              "product_id": "ETH-USD",
+              "price": "2601.76000000",
+              "remaining_size": "3.09000000",
+              "sequence": 2,
+              "time": "2017-06-25T11:23:14.775000Z"
+            })
+        ]
+        mock_connect.return_value.aenter.receive_str.side_effect = \
+            messages_expected
+        async with gdax.orderbook.OrderBook() as orderbook:
+            message = await orderbook.handle_message()
+            assert message == json.loads(messages_expected[0])
+
+            message = await orderbook.handle_message()
+            assert message is None
+
+            message = await orderbook.handle_message()
+            assert message == json.loads(messages_expected[2])
+
+    @patch('gdax.trader.Trader.get_product_order_book')
+    async def test_out_of_order(self, mock_book, mock_connect):
+        mock_connect.return_value.aenter.receive_str = CoroutineMock()
+        mock_connect.return_value.aenter.send_json = CoroutineMock()
+        mock_book.return_value = {'bids': [], 'asks': [], 'sequence': 1}
+
+        messages_expected = [
+            {
+              "type": "done",
+              "side": "sell",
+              "order_id": "4eef1226-4b38-422c-a5b1-56def7107f9a",
+              "reason": "canceled",
+              "product_id": "ETH-USD",
+              "price": "2601.76000000",
+              "remaining_size": "3.09000000",
+              "sequence": 2,
+              "time": "2017-06-25T11:23:14.775000Z"
+            },
+            {
+              "type": "done",
+              "side": "sell",
+              "order_id": "4eef1226-4b38-422c-a5b1-56def7107f9a",
+              "reason": "canceled",
+              "product_id": "ETH-USD",
+              "price": "2601.76000000",
+              "remaining_size": "3.09000000",
+              "sequence": 4,
+              "time": "2017-06-25T11:23:14.775000Z"
+            },
+            {
+              "type": "done",
+              "side": "sell",
+              "order_id": "4eef1226-4b38-422c-a5b1-56def7107f9a",
+              "reason": "canceled",
+              "product_id": "ETH-USD",
+              "price": "2601.76000000",
+              "remaining_size": "3.09000000",
+              "sequence": 2,
+              "time": "2017-06-25T11:23:14.775000Z"
+            },
+        ]
+        mock_connect.return_value.aenter.receive_str.side_effect = [
+            json.dumps(message_expected)
+            for message_expected in messages_expected
+        ]
+        async with gdax.orderbook.OrderBook() as orderbook:
+            message = await orderbook.handle_message()
+            assert message == messages_expected[0]
+
+            message = await orderbook.handle_message()
+            assert message is None
+
+            message = await orderbook.handle_message()
+            assert message == messages_expected[2]