Rework sandbox to clock

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

pkg/networkservice/chains/nsmgr/suite_test.go

@@ -42,6 +42,7 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/mechanisms/kernel"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/chain"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/utils/inject/injecterror"
+	"github.com/networkservicemesh/sdk/pkg/tools/clock"
 	"github.com/networkservicemesh/sdk/pkg/tools/sandbox"
 )
 
@@ -135,6 +136,8 @@ func (s *nsmgrSuite) Test_SelectsRestartingEndpointUsecase() {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
 
+	clockTime := clock.FromContext(ctx)
+
 	nsReg, err := s.domain.Nodes[0].NSRegistryClient.Register(ctx, defaultRegistryService())
 	require.NoError(t, err)
 
@@ -158,7 +161,7 @@ func (s *nsmgrSuite) Test_SelectsRestartingEndpointUsecase() {
 	// 2. Postpone endpoint start
 	time.AfterFunc(time.Second, func() {
 		serv := grpc.NewServer()
-		endpoint.NewServer(ctx, sandbox.GenerateTestToken).Register(serv)
+		endpoint.NewServer(ctx, sandbox.GenerateTestToken(clockTime)).Register(serv)
 		_ = serv.Serve(netListener)
 	})
 
@@ -387,6 +390,8 @@ func (s *nsmgrSuite) Test_PassThroughRemoteUsecase() {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
 
+	clockTime := clock.FromContext(ctx)
+
 	var nsesReg [nodesCount]*registry.NetworkServiceEndpoint
 	for i := 0; i < nodesCount; i++ {
 		var additionalFunctionality []networkservice.NetworkServiceServer
@@ -400,7 +405,7 @@ func (s *nsmgrSuite) Test_PassThroughRemoteUsecase() {
 							newPassTroughClient(fmt.Sprintf("my-service-remote-%v", i-1)),
 							kernel.NewClient()),
 						connect.WithDialTimeout(sandbox.DialTimeout),
-						connect.WithDialOptions(sandbox.DefaultDialOptions(sandbox.GenerateTestToken)...),
+						connect.WithDialOptions(sandbox.DefaultDialOptions(sandbox.GenerateTestToken(clockTime))...),
 					),
 				),
 			}
@@ -448,6 +453,8 @@ func (s *nsmgrSuite) Test_PassThroughLocalUsecase() {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
 
+	clockTime := clock.FromContext(ctx)
+
 	var nsesReg [nsesCount]*registry.NetworkServiceEndpoint
 	for i := 0; i < nsesCount; i++ {
 		var additionalFunctionality []networkservice.NetworkServiceServer
@@ -460,7 +467,7 @@ func (s *nsmgrSuite) Test_PassThroughLocalUsecase() {
 							newPassTroughClient(fmt.Sprintf("my-service-remote-%v", i-1)),
 							kernel.NewClient()),
 						connect.WithDialTimeout(sandbox.DialTimeout),
-						connect.WithDialOptions(sandbox.DefaultDialOptions(sandbox.GenerateTestToken)...),
+						connect.WithDialOptions(sandbox.DefaultDialOptions(sandbox.GenerateTestToken(clockTime))...),
 					),
 				),
 			}

pkg/networkservice/common/connect/server_cancel_test.go

@@ -41,6 +41,7 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/chain"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/next"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/utils/checks/checkrequest"
+	"github.com/networkservicemesh/sdk/pkg/tools/clock"
 	"github.com/networkservicemesh/sdk/pkg/tools/sandbox"
 )
 
@@ -50,6 +51,8 @@ func TestConnect_CancelDuringRequest(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
 
+	clockTime := clock.FromContext(ctx)
+
 	domain := sandbox.NewBuilder(ctx, t).
 		SetNodesCount(1).
 		SetRegistryProxySupplier(nil).
@@ -102,7 +105,7 @@ func TestConnect_CancelDuringRequest(t *testing.T) {
 			connect.NewServer(ctx,
 				clientFactory,
 				connect.WithDialTimeout(sandbox.DialTimeout),
-				connect.WithDialOptions(sandbox.DefaultDialOptions(sandbox.GenerateTestToken)...),
+				connect.WithDialOptions(sandbox.DefaultDialOptions(sandbox.GenerateTestToken(clockTime))...),
 			),
 		),
 	)

pkg/networkservice/common/heal/server_test.go

@@ -37,6 +37,7 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/chain"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/eventchannel"
 	"github.com/networkservicemesh/sdk/pkg/tools/addressof"
+	"github.com/networkservicemesh/sdk/pkg/tools/clock"
 	"github.com/networkservicemesh/sdk/pkg/tools/sandbox"
 )
 
@@ -71,18 +72,21 @@ func TestHealClient_Request(t *testing.T) {
 
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	defer cancelFunc()
+
+	clockTime := clock.FromContext(ctx)
+
 	monitorServer := eventchannel.NewMonitorServer(eventCh)
 	server := chain.NewNetworkServiceServer(
 		updatepath.NewServer("testServer"),
 		monitor.NewServer(ctx, &monitorServer),
-		updatetoken.NewServer(sandbox.GenerateTestToken),
+		updatetoken.NewServer(sandbox.GenerateTestToken(clockTime)),
 	)
 	healServer := heal.NewServer(ctx, addressof.NetworkServiceClient(onHeal))
 	client := chain.NewNetworkServiceClient(
 		updatepath.NewClient("testClient"),
 		adapters.NewServerToClient(healServer),
 		heal.NewClient(ctx, adapters.NewMonitorServerToClient(monitorServer)),
-		adapters.NewServerToClient(updatetoken.NewServer(sandbox.GenerateTestToken)),
+		adapters.NewServerToClient(updatetoken.NewServer(sandbox.GenerateTestToken(clockTime))),
 		adapters.NewServerToClient(server),
 	)
 

pkg/networkservice/common/refresh/client_test.go

@@ -212,11 +212,14 @@ func TestRefreshClient_CheckRaceConditions(t *testing.T) {
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
+
+	clockTime := clock.FromContext(ctx)
+
 	client := next.NewNetworkServiceClient(
 		serialize.NewClient(),
 		updatepath.NewClient("foo"),
 		refresh.NewClient(ctx),
-		adapters.NewServerToClient(updatetoken.NewServer(sandbox.GenerateExpiringToken(conf.expireTimeout))),
+		adapters.NewServerToClient(updatetoken.NewServer(sandbox.GenerateExpiringToken(conf.expireTimeout)(clockTime))),
 		adapters.NewServerToClient(refreshTester),
 	)
 
@@ -232,7 +235,6 @@ func TestRefreshClient_Sandbox(t *testing.T) {
 	domain := sandbox.NewBuilder(ctx, t).
 		SetNodesCount(2).
 		SetRegistryProxySupplier(nil).
-		SetTokenGenerateFunc(sandbox.GenerateTestToken).
 		Build()
 
 	nsReg := &registry.NetworkService{
@@ -250,8 +252,7 @@ func TestRefreshClient_Sandbox(t *testing.T) {
 	refreshSrv := newRefreshTesterServer(t, sandboxMinDuration, sandboxExpireTimeout)
 	domain.Nodes[0].NewEndpoint(ctx, nseReg, sandbox.GenerateTestToken, refreshSrv)
 
-	nscTokenGenerator := sandbox.GenerateExpiringToken(sandboxExpireTimeout)
-	nsc := domain.Nodes[1].NewClient(ctx, nscTokenGenerator)
+	nsc := domain.Nodes[1].NewClient(ctx, sandbox.GenerateExpiringToken(sandboxExpireTimeout))
 
 	refreshSrv.beforeRequest("test-conn")
 	_, err = nsc.Request(ctx, mkRequest("test-conn", nil))

pkg/registry/chains/memory/server_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/networkservicemesh/api/pkg/api/networkservice/payload"
 	"github.com/networkservicemesh/api/pkg/api/registry"
 
+	"github.com/networkservicemesh/sdk/pkg/tools/clock"
 	"github.com/networkservicemesh/sdk/pkg/tools/grpcutils"
 	"github.com/networkservicemesh/sdk/pkg/tools/sandbox"
 
@@ -41,14 +42,16 @@ func Test_RegistryMemory_ShouldSetDefaultPayload(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
 
+	clockTime := clock.FromContext(ctx)
+
 	domain := sandbox.NewBuilder(ctx, t).
 		SetNodesCount(0).
 		SetRegistryProxySupplier(nil).
 		SetNSMgrProxySupplier(nil).
 		Build()
 
 	// start grpc client connection and register it
-	cc, err := grpc.DialContext(ctx, grpcutils.URLToTarget(domain.Registry.URL), sandbox.DefaultDialOptions(sandbox.GenerateTestToken)...)
+	cc, err := grpc.DialContext(ctx, grpcutils.URLToTarget(domain.Registry.URL), sandbox.DefaultDialOptions(sandbox.GenerateTestToken(clockTime))...)
 	require.NoError(t, err)
 	defer func() {
 		_ = cc.Close()

pkg/tools/grpcutils/server_stream.go

@@ -0,0 +1,43 @@
+// Copyright (c) 2021 Doc.ai and/or its affiliates.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at:
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package grpcutils
+
+import (
+	"context"
+
+	"google.golang.org/grpc"
+)
+
+// ServerStream is a grpc.ServerStream wrapper with changed context
+type ServerStream struct {
+	ctx context.Context
+
+	grpc.ServerStream
+}
+
+// Context return ss.ctx
+func (ss *ServerStream) Context() context.Context {
+	return ss.ctx
+}
+
+// WrapServerStreamContext wraps given grpc.ServerStream into ServerStream with ctx = wrapper(ss.Context())
+func WrapServerStreamContext(ss grpc.ServerStream, wrapper func(ctx context.Context) context.Context) *ServerStream {
+	return &ServerStream{
+		ctx:          wrapper(ss.Context()),
+		ServerStream: ss,
+	}
+}