Refactor tests

Signed-off-by: Vladimir Popov <vladimir.popov@xored.com>

pkg/networkservice/chains/nsmgr/server_test.go

@@ -548,39 +548,41 @@ func testNSEAndClient(
 	domain *sandbox.Domain,
 	nseReg *registry.NetworkServiceEndpoint,
 ) {
+	// 1. Create an endpoint
 	nseCtx, nseCancel := context.WithCancel(ctx)
 	defer nseCancel()
 
 	_, err := sandbox.NewEndpoint(nseCtx, nseReg, sandbox.GenerateTestToken, domain.Nodes[0].NSMgr)
 	require.NoError(t, err)
 
+	// 2. Create a client
 	clientCtx, clientCancel := context.WithCancel(ctx)
 	defer clientCancel()
 
 	nsc := sandbox.NewClient(clientCtx, sandbox.GenerateTestToken, domain.Nodes[0].NSMgr.URL)
 
-	conn, err := func() (*networkservice.Connection, error) {
-		requestCtx, requestCancel := context.WithCancel(ctx)
-		defer requestCancel()
-
-		return nsc.Request(requestCtx, &networkservice.NetworkServiceRequest{
-			MechanismPreferences: []*networkservice.Mechanism{
-				{Cls: cls.LOCAL, Type: kernelmech.MECHANISM},
-			},
-			Connection: &networkservice.Connection{
-				NetworkService: nseReg.NetworkServiceNames[0],
-			},
-		})
-	}()
+	// 3. Request endpoint
+	requestCtx, requestCancel := context.WithCancel(ctx)
+
+	conn, err := nsc.Request(requestCtx, &networkservice.NetworkServiceRequest{
+		MechanismPreferences: []*networkservice.Mechanism{
+			{Cls: cls.LOCAL, Type: kernelmech.MECHANISM},
+		},
+		Connection: &networkservice.Connection{
+			NetworkService: nseReg.NetworkServiceNames[0],
+		},
+	})
 	require.NoError(t, err)
 
-	_, err = func() (interface{}, error) {
-		closeCtx, closeCancel := context.WithCancel(ctx)
-		defer closeCancel()
+	requestCancel()
 
-		return nsc.Close(closeCtx, conn)
-	}()
+	// 4. Close connection
+	closeCtx, closeCancel := context.WithCancel(ctx)
+
+	_, err = nsc.Close(closeCtx, conn)
 	require.NoError(t, err)
+
+	closeCancel()
 }
 
 type passThroughClient struct {

pkg/registry/common/querycache/nse_client_test.go

@@ -132,26 +132,26 @@ func Test_QueryCacheServer_ShouldCleanupGoroutinesOnNSEUnregister(t *testing.T)
 
 	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
 
-	_, err = func() (interface{}, error) {
-		findCtx, findCancel := context.WithCancel(ctx)
-		defer findCancel()
-
-		return client.Find(findCtx, &registry.NetworkServiceEndpointQuery{
-			NetworkServiceEndpoint: &registry.NetworkServiceEndpoint{
-				Name: reg.Name,
-			},
-		})
-	}()
+	// 1. Find
+	findCtx, findCancel := context.WithCancel(ctx)
+
+	_, err = client.Find(findCtx, &registry.NetworkServiceEndpointQuery{
+		NetworkServiceEndpoint: &registry.NetworkServiceEndpoint{
+			Name: reg.Name,
+		},
+	})
 	require.NoError(t, err)
 
-	// Wait a bit for the (cache -> registry) stream to start
+	findCancel()
+
+	// 2. Wait a bit for the (cache -> registry) stream to start
 	<-time.After(1 * time.Millisecond)
 
-	_, err = func() (interface{}, error) {
-		unregisterCtx, unregisterCancel := context.WithCancel(ctx)
-		defer unregisterCancel()
+	// 3. Unregister
+	unregisterCtx, unregisterCancel := context.WithCancel(ctx)
 
-		return mem.Unregister(unregisterCtx, reg)
-	}()
+	_, err = mem.Unregister(unregisterCtx, reg)
 	require.NoError(t, err)
+
+	unregisterCancel()
 }