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