genericWebsocket: add reconnect attempt on disconnect

bfxapi/websockets/GenericWebsocket.py

@@ -4,11 +4,14 @@ Module used as a interfeace to describe a generick websocket client
 
 import asyncio
 import websockets
+import socket
 import json
 
 from pyee import EventEmitter
 from ..utils.CustomLogger import CustomLogger
 
+# websocket exceptions
+from websockets.exceptions import ConnectionClosed
 
 class AuthError(Exception):
     """
@@ -16,7 +19,6 @@ class AuthError(Exception):
     """
     pass
 
-
 def is_json(myjson):
     try:
         json_object = json.loads(myjson)
@@ -24,20 +26,20 @@ def is_json(myjson):
         return False
     return True
 
-
 class GenericWebsocket:
     """
     Websocket object used to contain the base functionality of a websocket.
     Inlcudes an event emitter and a standard websocket client.
     """
 
-    def __init__(self, host, logLevel='INFO', loop=None):
+    def __init__(self, host, logLevel='INFO', loop=None, max_retries=5):
         self.host = host
         self.logger = CustomLogger('BfxWebsocket', logLevel=logLevel)
         self.loop = loop or asyncio.get_event_loop()
         self.events = EventEmitter(
             scheduler=asyncio.ensure_future, loop=self.loop)
         self.ws = None
+        self.max_retries = max_retries
 
     def run(self):
         """
@@ -51,15 +53,33 @@ class GenericWebsocket:
         """
         return self._main(self.host)
 
-    async def _main(self, host):
+    async def _connect(self, host):
         async with websockets.connect(host) as websocket:
             self.ws = websocket
-            self.logger.info("Wesocket connectedt to {}".format(self.host))
+            self.logger.info("Wesocket connected to {}".format(self.host))
             while True:
                 await asyncio.sleep(0)
                 message = await websocket.recv()
                 await self.on_message(message)
 
+    def get_ws(self):
+        return self.ws
+
+    async def _main(self, host):
+        retries = 0
+        while retries < self.max_retries:
+            try:
+                await self._connect(host)
+                retries = 0
+            except (ConnectionClosed, socket.error) as e:
+                self.logger.error(str(e))
+                retries += 1
+                # wait 5 seconds befor retrying
+                self.logger.info("Waiting 5 seconds befor retrying...")
+                await asyncio.sleep(5)
+                self.logger.info("Reconnect attempt {}/{}".format(retries, self.max_retries))
+        self.logger.info("Unable to connect to websocket.")
+
     def remove_all_listeners(self, event):
         """
         Remove all listeners from event emitter