diff --git a/polygon/websocket/websocket_client.py b/polygon/websocket/websocket_client.py
index b42515f3..3c046485 100644
--- a/polygon/websocket/websocket_client.py
+++ b/polygon/websocket/websocket_client.py
@@ -17,7 +17,8 @@ class WebSocketClient:
     #  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,
-                 on_close: Optional[Callable[[str], None]] = None, on_error: Optional[Callable[[str], None]] = None):
+                 on_close: Optional[Callable[[websocket.WebSocketApp], None]] = None,
+                 on_error: Optional[Callable[[websocket.WebSocketApp, 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(),
diff --git a/websocket-example.py b/websocket-example.py
index 8218a4b6..10ce55a2 100644
--- a/websocket-example.py
+++ b/websocket-example.py
@@ -1,15 +1,23 @@
 import time
 
-from polygon_client import WebSocketClient, STOCKS_CLUSTER
+from polygon import WebSocketClient, STOCKS_CLUSTER
 
 
-def my_customer_process_message(message):
+def my_custom_process_message(message):
     print("this is my custom message processing", message)
 
 
+def my_custom_error_handler(ws, error):
+    print("this is my custom error handler", error)
+
+
+def my_custom_close_handler(ws):
+    print("this is my custom close handler")
+
+
 def main():
     key = 'your api key'
-    my_client = WebSocketClient(STOCKS_CLUSTER, key, my_customer_process_message)
+    my_client = WebSocketClient(STOCKS_CLUSTER, key, my_custom_process_message)
     my_client.run_async()
 
     my_client.subscribe("T.MSFT", "T.AAPL", "T.AMD", "T.NVDA")