From d281ae2b7afad7bb8d1f405befa142959d0f6c0c Mon Sep 17 00:00:00 2001
From: Vladimir Popov <vladimir.popov@xored.com>
Date: Fri, 29 Jan 2021 18:39:03 +0700
Subject: [PATCH] Add sandbox.Unregister()

Signed-off-by: Vladimir Popov <vladimir.popov@xored.com>
---
 .../chains/nsmgr/server_test.go               | 30 +++++----------
 pkg/tools/sandbox/utils.go                    | 37 ++++++++++++-------
 2 files changed, 33 insertions(+), 34 deletions(-)

diff --git a/pkg/networkservice/chains/nsmgr/server_test.go b/pkg/networkservice/chains/nsmgr/server_test.go
index c16fc78c10..23f9b79073 100644
--- a/pkg/networkservice/chains/nsmgr/server_test.go
+++ b/pkg/networkservice/chains/nsmgr/server_test.go
@@ -510,7 +510,7 @@ func TestNSMGR_PassThroughLocal(t *testing.T) {
 func TestNSMGR_ShouldCleanAllClientAndEndpointGoroutines(t *testing.T) {
 	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5000)
 	defer cancel()
 
 	domain := sandbox.NewBuilder(t).
@@ -548,23 +548,15 @@ func testNSEAndClient(
 	domain *sandbox.Domain,
 	nseReg *registry.NetworkServiceEndpoint,
 ) {
-	// 1. Create an endpoint
-	nseCtx, nseCancel := context.WithCancel(ctx)
-	defer nseCancel()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 
-	_, err := sandbox.NewEndpoint(nseCtx, nseReg, sandbox.GenerateTestToken, domain.Nodes[0].NSMgr)
+	_, err := sandbox.NewEndpoint(ctx, 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)
-
-	// 3. Request endpoint
-	requestCtx, requestCancel := context.WithCancel(ctx)
+	nsc := sandbox.NewClient(ctx, sandbox.GenerateTestToken, domain.Nodes[0].NSMgr.URL)
 
-	conn, err := nsc.Request(requestCtx, &networkservice.NetworkServiceRequest{
+	conn, err := nsc.Request(ctx, &networkservice.NetworkServiceRequest{
 		MechanismPreferences: []*networkservice.Mechanism{
 			{Cls: cls.LOCAL, Type: kernelmech.MECHANISM},
 		},
@@ -574,15 +566,11 @@ func testNSEAndClient(
 	})
 	require.NoError(t, err)
 
-	requestCancel()
-
-	// 4. Close connection
-	closeCtx, closeCancel := context.WithCancel(ctx)
-
-	_, err = nsc.Close(closeCtx, conn)
+	_, err = nsc.Close(ctx, conn)
 	require.NoError(t, err)
 
-	closeCancel()
+	err = sandbox.UnregisterEndpoint(ctx, nseReg, domain.Nodes[0].NSMgr)
+	require.NoError(t, err)
 }
 
 type passThroughClient struct {
diff --git a/pkg/tools/sandbox/utils.go b/pkg/tools/sandbox/utils.go
index 16a584e566..e608d56d5a 100644
--- a/pkg/tools/sandbox/utils.go
+++ b/pkg/tools/sandbox/utils.go
@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"github.com/networkservicemesh/api/pkg/api/networkservice/payload"
+	"github.com/pkg/errors"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/networkservicemesh/sdk/pkg/tools/logger"
@@ -64,17 +65,22 @@ func NewEndpoint(ctx context.Context, nse *registry.NetworkServiceEndpoint, gene
 		}
 	}
 	serve(ctx, u, ep.Register)
+
 	if nse.Url == "" {
 		nse.Url = u.String()
 	}
-
 	if nse.ExpirationTime == nil {
 		nse.ExpirationTime = timestamppb.New(time.Now().Add(time.Hour))
 	}
-	if nse, err = mgr.NetworkServiceEndpointRegistryServer().Register(ctx, nse); err != nil {
+
+	var reg *registry.NetworkServiceEndpoint
+	if reg, err = mgr.NetworkServiceEndpointRegistryServer().Register(ctx, nse); err != nil {
 		return nil, err
 	}
 
+	nse.Name = reg.Name
+	nse.ExpirationTime = reg.ExpirationTime
+
 	for _, service := range nse.NetworkServiceNames {
 		if _, err := mgr.NetworkServiceRegistryServer().Register(ctx, &registry.NetworkService{
 			Name:    service,
@@ -84,22 +90,27 @@ func NewEndpoint(ctx context.Context, nse *registry.NetworkServiceEndpoint, gene
 		}
 	}
 
-	go func() {
-		<-ctx.Done()
-		for _, service := range nse.NetworkServiceNames {
-			_, _ = mgr.NetworkServiceRegistryServer().Unregister(context.Background(), &registry.NetworkService{
-				Name:    service,
-				Payload: payload.IP,
-			})
-		}
-		_, _ = mgr.NetworkServiceEndpointRegistryServer().Unregister(context.Background(), nse)
-	}()
-
 	logger.Log(ctx).Infof("Started listen endpoint %v on %v.", nse.Name, u.String())
 
 	return &EndpointEntry{Endpoint: ep, URL: u}, nil
 }
 
+// UnregisterEndpoint unregisters endpoint from NSMgr
+func UnregisterEndpoint(ctx context.Context, nse *registry.NetworkServiceEndpoint, mgr nsmgr.Nsmgr) (err error) {
+	if _, unregisterErr := mgr.NetworkServiceEndpointRegistryServer().Unregister(ctx, nse); unregisterErr != nil {
+		err = errors.Wrapf(unregisterErr, "%v\n", err)
+	}
+	for _, name := range nse.NetworkServiceNames {
+		if _, unregisterErr := mgr.NetworkServiceRegistryServer().Unregister(ctx, &registry.NetworkService{
+			Name:    name,
+			Payload: payload.IP,
+		}); unregisterErr != nil {
+			err = errors.Wrapf(unregisterErr, "%v\n", err)
+		}
+	}
+	return err
+}
+
 // NewClient is a client.NewClient over *url.URL with some fields preset for testing
 func NewClient(ctx context.Context, generatorFunc token.GeneratorFunc, connectTo *url.URL, additionalFunctionality ...networkservice.NetworkServiceClient) networkservice.NetworkServiceClient {
 	return clienturl.NewClient(