From 3332c985bd9739d189ddd93eb65b3096b6e2012a Mon Sep 17 00:00:00 2001
From: jbonzo <8647805+jbonzo@users.noreply.github.com>
Date: Fri, 31 Jan 2020 12:28:55 -0500
Subject: [PATCH] Add on_error on_close

---
 polygon/websocket/websocket_client.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/polygon/websocket/websocket_client.py b/polygon/websocket/websocket_client.py
index cb750cbe..b42515f3 100644
--- a/polygon/websocket/websocket_client.py
+++ b/polygon/websocket/websocket_client.py
@@ -16,7 +16,8 @@ class WebSocketClient:
     #  the 3 possible clusters (I think I like client per, but then a problem is the user can make multiple clients for
     #  the same cluster and that's not desirable behavior,
     #  somehow keeping track with multiple Client instances will be the difficulty)
-    def __init__(self, cluster: str, auth_key: str, process_message: Optional[Callable[[str], None]] = None):
+    def __init__(self, cluster: str, auth_key: str, process_message: Optional[Callable[[str], None]] = None,
+                 on_close: Optional[Callable[[str], None]] = None, on_error: Optional[Callable[[str], None]] = None):
         self._host = self.DEFAULT_HOST
         self.url = f"wss://{self._host}/{cluster}"
         self.ws: websocket.WebSocketApp = websocket.WebSocketApp(self.url, on_open=self._default_on_open(),
@@ -26,6 +27,8 @@ def __init__(self, cluster: str, auth_key: str, process_message: Optional[Callab
         self.auth_key = auth_key
 
         self.process_message = process_message
+        self.ws.on_close = on_close
+        self.ws.on_error = on_error
 
         # being authenticated is an event that must occur before any other action is sent to the server
         self._authenticated = threading.Event()