Fixed monitor tests

Signed-off-by: Sergey Semenov <[email protected]>
networkservicemesh · Mar 17, 2020 · 42a1051 · 42a1051
1 parent b1dddf0
commit 42a1051
Showing 1 changed file with 33 additions and 18 deletions.
diff --git a/pkg/networkservice/common/monitor/server_test.go b/pkg/networkservice/common/monitor/server_test.go
@@ -32,8 +32,22 @@ import (
 
 func TestMonitor(t *testing.T) {
 	defer goleak.VerifyNone(t)
+
 	// Specify pathSegments to test
-	segmentNames := []string{"local-nsm", "remote-nsm"}
+	ctx1, cancelFunc1 := context.WithCancel(context.Background())
+	ctx2, cancelFunc2 := context.WithCancel(context.Background())
+	segmentData := []struct {
+		name       string
+		ctx        context.Context
+		cancelFunc context.CancelFunc
+	}{
+		{
+			"local-nsm", ctx1, cancelFunc1,
+		},
+		{
+			"remote-nsm", ctx2, cancelFunc2,
+		},
+	}
 
 	// Create monitorServer, monitorClient, and server.
 	var monitorServer networkservice.MonitorConnectionServer
@@ -45,30 +59,30 @@ func TestMonitor(t *testing.T) {
 	receivers := make(map[string]networkservice.MonitorConnection_MonitorConnectionsClient)
 
 	// Get Empty initial state transfers
-	for _, segmentName := range segmentNames {
+	for _, segment := range segmentData {
 		var err error
-		receivers[segmentName], err = monitorClient.MonitorConnections(context.Background(), &networkservice.MonitorScopeSelector{
-			PathSegments: []*networkservice.PathSegment{{Name: segmentName}},
+		receivers[segment.name], err = monitorClient.MonitorConnections(segment.ctx, &networkservice.MonitorScopeSelector{
+			PathSegments: []*networkservice.PathSegment{{Name: segment.name}},
 		})
 		assert.Nil(t, err)
-		event, err := receivers[segmentName].Recv()
+		event, err := receivers[segment.name].Recv()
 		assert.Nil(t, err)
 		require.NotNil(t, event)
 		assert.Equal(t, networkservice.ConnectionEventType_INITIAL_STATE_TRANSFER, event.GetType())
-		require.Equal(t, len(event.GetConnections()[segmentName].GetPath().GetPathSegments()), 0)
+		require.Equal(t, len(event.GetConnections()[segment.name].GetPath().GetPathSegments()), 0)
 	}
 
 	// Send requests
-	for _, segmentName := range segmentNames {
+	for _, segment := range segmentData {
 		var err error
-		connections[segmentName], err = server.Request(context.Background(), &networkservice.NetworkServiceRequest{
+		connections[segment.name], err = server.Request(context.Background(), &networkservice.NetworkServiceRequest{
 			Connection: &networkservice.Connection{
-				Id: segmentName,
+				Id: segment.name,
 				Path: &networkservice.Path{
 					Index: 0,
 					PathSegments: []*networkservice.PathSegment{
 						{
-							Name: segmentName,
+							Name: segment.name,
 						},
 					},
 				},
@@ -78,13 +92,13 @@ func TestMonitor(t *testing.T) {
 	}
 
 	// Get Updates and insure we've properly filtered by segmentName
-	for _, segmentName := range segmentNames {
-		event, err := receivers[segmentName].Recv()
+	for _, segment := range segmentData {
+		event, err := receivers[segment.name].Recv()
 		assert.Nil(t, err)
 		require.NotNil(t, event)
 		assert.Equal(t, networkservice.ConnectionEventType_UPDATE, event.GetType())
-		require.Equal(t, len(event.GetConnections()[segmentName].GetPath().GetPathSegments()), 1)
-		assert.Equal(t, segmentName, event.GetConnections()[segmentName].GetPath().GetPathSegments()[0].GetName())
+		require.Equal(t, len(event.GetConnections()[segment.name].GetPath().GetPathSegments()), 1)
+		assert.Equal(t, segment.name, event.GetConnections()[segment.name].GetPath().GetPathSegments()[0].GetName())
 	}
 
 	// Close Connections
@@ -94,12 +108,13 @@ func TestMonitor(t *testing.T) {
 	}
 
 	// Get Delete Events and insure we've properly filtered by segmentName
-	for _, segmentName := range segmentNames {
-		event, err := receivers[segmentName].Recv()
+	for _, segment := range segmentData {
+		event, err := receivers[segment.name].Recv()
 		assert.Nil(t, err)
 		require.NotNil(t, event)
 		assert.Equal(t, networkservice.ConnectionEventType_DELETE, event.GetType())
-		require.Equal(t, len(event.GetConnections()[segmentName].GetPath().GetPathSegments()), 1)
-		assert.Equal(t, segmentName, event.GetConnections()[segmentName].GetPath().GetPathSegments()[0].GetName())
+		require.Equal(t, len(event.GetConnections()[segment.name].GetPath().GetPathSegments()), 1)
+		assert.Equal(t, segment.name, event.GetConnections()[segment.name].GetPath().GetPathSegments()[0].GetName())
+		segment.cancelFunc()
 	}
 }