improve orderbook tests

Author: csko

gdax/orderbook.py

@@ -20,6 +20,10 @@
 import gdax.utils
 
 
+class OrderBookError(Exception):
+    pass
+
+
 class OrderBook(object):
     def __init__(self, product_ids='ETH-USD', api_key=None, api_secret=None,
                  passphrase=None, use_heartbeat=False,
@@ -147,6 +151,11 @@ async def handle_message(self):
             await self._init()
             return
 
+        msg_type = message['type']
+
+        if msg_type == 'error':
+            raise OrderBookError(f'Error: {message["message"]}')
+
         product_id = message['product_id']
         assert self._sequences[product_id] is not None
         sequence = message['sequence']
@@ -163,7 +172,6 @@ async def handle_message(self):
             await self._init()
             return
 
-        msg_type = message['type']
         if msg_type == 'open':
             self.add(product_id, message)
         elif msg_type == 'done' and 'price' in message:
@@ -172,16 +180,14 @@ async def handle_message(self):
             self.match(product_id, message)
         elif msg_type == 'change':
             self.change(product_id, message)
-        elif msg_type == 'error':
-            raise Exception(f'Error: {message["message"]}')
         elif msg_type == 'heartbeat':
             pass
         elif msg_type == 'received':
             pass
         elif msg_type == 'done':
             pass
         else:
-            raise Exception(f'unknown message type {msg_type}')
+            raise OrderBookError(f'unknown message type {msg_type}')
 
         self._sequences[product_id] = sequence
         return message
@@ -344,7 +350,7 @@ def get_max_bid_depth(self, product_id):
         return sum([order['size'] for order in orders])
 
 
-async def run_orderbook():
+async def run_orderbook():  # pragma: no cover
     async with OrderBook(
         ['ETH-USD', 'BTC-USD'],
         api_key=None,
@@ -367,7 +373,7 @@ async def run_orderbook():
                          orderbook.get_min_ask_depth('ETH-USD'),
                          orderbook.get_max_bid_depth('ETH-USD'))
 
-if __name__ == "__main__":
+if __name__ == "__main__":  # pragma: no cover
     logging.getLogger().setLevel(logging.INFO)
     loop = asyncio.get_event_loop()
     loop.run_until_complete(run_orderbook())

gdax/trader.py

@@ -349,7 +349,7 @@ async def get_trailing_volume(self):
         return await self._get("/users/self/trailing-volume")
 
 
-async def main():
+async def main():  # pragma: no cover
     trader = Trader(product_id='BTC-USD',
                     api_key=None,
                     api_secret=None,
@@ -363,7 +363,7 @@ async def main():
     logging.info(res)
 
 
-if __name__ == "__main__":
+if __name__ == "__main__":  # pragma: no cover
     logging.getLogger().setLevel(logging.INFO)
     loop = asyncio.get_event_loop()
     loop.run_until_complete(main())

tests/test_orderbook.py

@@ -304,6 +304,7 @@ async def test_orderbook_advanced(self, mock_book, mock_connect):
         mock_connect.return_value.aenter.send_json = CoroutineMock()
         messages_expected = [
             {  # ignored
+              "type": "received",
               "product_id": product_id,
               "sequence": sequence - 1,
             },
@@ -454,3 +455,30 @@ async def test_orderbook_advanced(self, mock_book, mock_connect):
             current_book['sequence'] += 1
             assert orderbook.get_current_book(product_id) != current_book
             # TODO
+
+    @patch('gdax.trader.Trader.get_product_order_book')
+    async def test_error(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": "error",
+            "message": "test error",
+          },
+          {
+            "type": "unknownmsgtype",
+            'product_id': 'ETH-USD',
+            'sequence': 2,
+          },
+        ]
+        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:
+            with pytest.raises(gdax.orderbook.OrderBookError):
+                message = await orderbook.handle_message()
+
+            with pytest.raises(gdax.orderbook.OrderBookError):
+                message = await orderbook.handle_message()