Add context for Watch.

Author: cloudwebrtc

pkg/client/client.go

@@ -113,7 +113,7 @@ func (c *Client) handleNatsMsg(msg *nats.Msg, callback NodeStateChangeCallback)
 	return nil
 }
 
-func (c *Client) Watch(service string, handleNodeState NodeStateChangeCallback) error {
+func (c *Client) Watch(ctx context.Context, service string, handleNodeState NodeStateChangeCallback) error {
 	if handleNodeState == nil {
 		err := fmt.Errorf("Watch callback must be set for %v", service)
 		logger.Warnf("Watch: err => %v", err)
@@ -138,6 +138,8 @@ func (c *Client) Watch(service string, handleNodeState NodeStateChangeCallback)
 			select {
 			case <-c.ctx.Done():
 				return c.ctx.Err()
+			case <-ctx.Done():
+				return ctx.Err()
 			case msg, ok := <-msgCh:
 				if ok {
 					err := c.handleNatsMsg(msg, handleNodeState)

pkg/client/client_test.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"context"
 	"sync"
 	"testing"
 
@@ -51,7 +52,7 @@ func TestWatch(t *testing.T) {
 		ExtraInfo: extraInfo,
 	}
 
-	s.Watch("sfu", func(state discovery.NodeState, n *discovery.Node) {
+	s.Watch(context.Background(), "sfu", func(state discovery.NodeState, n *discovery.Node) {
 		if state == discovery.NodeUp {
 			log.Infof("NodeUp => %v", *n)
 			assert.Equal(t, node, *n)