add more unit tests

Signed-off-by: NikitaSkrynnik <[email protected]>

extendtimeout/connection.go

@@ -35,6 +35,10 @@ type extendedContext struct {
 	valuesContext context.Context
 }
 
+func (ec *extendedContext) Value(key interface{}) interface{} {
+	return ec.valuesContext.Value(key)
+}
+
 // NewConnection - creates a wrapper for vpp connection that uses extended context timeout for all operations
 func NewConnection(vppConn api.Connection, contextTimeout time.Duration) api.Connection {
 	return &extendedConnection{

extendtimeout/connection_test.go

@@ -18,51 +18,103 @@ package extendtimeout_test
 
 import (
 	"context"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
 	"go.fd.io/govpp/api"
+	"go.uber.org/goleak"
 
 	"github.com/networkservicemesh/vpphelper/extendtimeout"
 )
 
 type testConn struct {
 	api.Connection
-	invoke func(ctx context.Context)
+	invokeBody func(ctx context.Context)
 }
 
 func (c *testConn) Invoke(ctx context.Context, req, reply api.Message) error {
-	c.invoke(ctx)
+	c.invokeBody(ctx)
 	return nil
 }
 
+type key struct{}
+
+const value = "value"
+
 func TestSmallTimeout(t *testing.T) {
-	testConn := &testConn{invoke: func(ctx context.Context) {
+	testConn := &testConn{invokeBody: func(ctx context.Context) {
 		deadline, ok := ctx.Deadline()
 		require.True(t, ok)
 		timeout := time.Until(deadline)
 		require.Greater(t, timeout, time.Second)
+		require.Equal(t, ctx.Value(&key{}), value)
 	}}
 
 	smallCtx, cancel := context.WithTimeout(context.Background(), time.Millisecond*10)
+	smallCtx = context.WithValue(smallCtx, &key{}, value)
 	defer cancel()
 
 	err := extendtimeout.NewConnection(testConn, 2*time.Second).Invoke(smallCtx, nil, nil)
 	require.NoError(t, err)
 }
 
 func TestBigTimeout(t *testing.T) {
-	testConn := &testConn{invoke: func(ctx context.Context) {
+	testConn := &testConn{invokeBody: func(ctx context.Context) {
 		deadline, ok := ctx.Deadline()
 		require.True(t, ok)
 		timeout := time.Until(deadline)
 		require.Greater(t, timeout, 7*time.Second)
+		require.Equal(t, ctx.Value(&key{}), value)
 	}}
 
 	bigCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	bigCtx = context.WithValue(bigCtx, &key{}, value)
 	defer cancel()
 
 	err := extendtimeout.NewConnection(testConn, 2*time.Second).Invoke(bigCtx, nil, nil)
 	require.NoError(t, err)
 }
+
+func TestNoTimeout(t *testing.T) {
+	testConn := &testConn{invokeBody: func(ctx context.Context) {
+		_, ok := ctx.Deadline()
+		require.False(t, ok)
+	}}
+
+	emptyCtx := context.Background()
+	err := extendtimeout.NewConnection(testConn, 2*time.Second).Invoke(emptyCtx, nil, nil)
+	require.NoError(t, err)
+}
+
+func TestCanceledContext(t *testing.T) {
+	t.Cleanup(func() {
+		goleak.VerifyNone(t)
+	})
+
+	counter := new(atomic.Int32)
+	ch := make(chan struct{}, 1)
+	defer close(ch)
+	testConn := &testConn{invokeBody: func(ctx context.Context) {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ch:
+			counter.Add(1)
+		}
+	}}
+
+	cancelCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	go func() {
+		err := extendtimeout.NewConnection(testConn, 20*time.Second).Invoke(cancelCtx, nil, nil)
+		require.NoError(t, err)
+	}()
+
+	cancel()
+	ch <- struct{}{}
+
+	require.Eventually(t, func() bool {
+		return counter.Load() == 1
+	}, time.Second, 100*time.Millisecond)
+}

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.4
 	go.fd.io/govpp v0.10.0-alpha.0.20240110141843-761adec77524
+	go.uber.org/goleak v1.3.0
 	gopkg.in/fsnotify.v1 v1.4.7
 )
 

go.sum

@@ -35,6 +35,8 @@ github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcU
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 go.fd.io/govpp v0.10.0-alpha.0.20240110141843-761adec77524 h1:Dja0i7Ln/aUAc10VN4Pp3SibYTfNKS1CPkM5TSGL5jk=
 go.fd.io/govpp v0.10.0-alpha.0.20240110141843-761adec77524/go.mod h1:9QoqjEbvfuuXNfjHS0A7YS+7QQVVaQ9cMioOWpSM4rY=
+go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
+go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
 golang.org/x/net v0.21.0 h1:AQyQV4dYCvJ7vGmJyKki9+PBdyvhkSd8EIx/qb0AYv4=
 golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=