From 4487f18e838a347102e764d0585493709d93568a Mon Sep 17 00:00:00 2001
From: Elara Musayelyan <elara@elara.ws>
Date: Mon, 12 Dec 2022 18:36:56 -0800
Subject: [PATCH] Add support for connection handlers

---
 websocket.go | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/websocket.go b/websocket.go
index e3af8e1..a747e28 100644
--- a/websocket.go
+++ b/websocket.go
@@ -59,6 +59,19 @@ func NewWebSocket(baseURL string) (*WSClient, error) {
 	return out, nil
 }
 
+// SetConnectHandler sets the connection handler function
+func (c *WSClient) SetConnectHandler(f func() error) {
+	c.conn.SubscribeHandler = f
+}
+
+// ConnectHandler invokes the connection handler function
+func (c *WSClient) ConnectHandler() error {
+	if c.conn.SubscribeHandler == nil {
+		return nil
+	}
+	return c.conn.SubscribeHandler()
+}
+
 func (c *WSClient) Login(ctx context.Context, l types.Login) error {
 	u := &url.URL{}
 	*u = *c.baseURL