Add TimeoutWithinContext to mongoutil

Author: prestonvasquez

internal/mongoutil/mongoutil.go

@@ -7,7 +7,9 @@
 package mongoutil
 
 import (
+	"context"
 	"reflect"
+	"time"
 
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
@@ -83,3 +85,17 @@ func HostsFromURI(uri string) ([]string, error) {
 
 	return opts.Hosts, nil
 }
+
+// TimeoutWithinContext will return true if the provided timeout is nil or if
+// it is less than the context deadline. If the context does not have a
+// deadline, it will return true.
+func TimeoutWithinContext(ctx context.Context, timeout time.Duration) bool {
+	deadline, ok := ctx.Deadline()
+	if !ok {
+		return true
+	}
+
+	ctxTimeout := time.Until(deadline)
+
+	return ctxTimeout <= 0 || timeout < ctxTimeout
+}

internal/mongoutil/mongoutil_test.go

@@ -7,9 +7,13 @@
 package mongoutil
 
 import (
+	"context"
 	"strings"
 	"testing"
+	"time"
 
+	"go.mongodb.org/mongo-driver/v2/internal/assert"
+	"go.mongodb.org/mongo-driver/v2/internal/ptrutil"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
 
@@ -32,3 +36,77 @@ func BenchmarkNewOptions(b *testing.B) {
 		}
 	})
 }
+
+func TestValidChangeStreamTimeouts(t *testing.T) {
+	tests := []struct {
+		name        string
+		ctxTimeout  *time.Duration
+		timeout     time.Duration
+		wantTimeout time.Duration
+		want        bool
+	}{
+		{
+			name:       "No timeout provided, no context deadline",
+			ctxTimeout: nil,
+			timeout:    nil,
+			want:       true,
+		},
+		{
+			name:       "Timeout shorter than context deadline",
+			ctxTimeout: ptrutil.Ptr(10 * time.Second),
+			timeout:    1 * time.Second,
+			want:       true,
+		},
+		{
+			name:       "Timeout equal to context deadline",
+			ctxTimeout: ptrutil.Ptr(1 * time.Second),
+			timeout:    1 * time.Second,
+			want:       false,
+		},
+		{
+			name:       "Timeout greater than context deadline",
+			ctxTimeout: ptrutil.Ptr(1 * time.Second),
+			timeout:    10 * time.Second,
+			want:       false,
+		},
+		{
+			name:       "Context deadline already expired",
+			ctxTimeout: ptrutil.Ptr(-1 * time.Second),
+			timeout:    1 * time.Second,
+			want:       true, // *timeout <= 0 branch in code
+		},
+		{
+			name:       "Timeout is zero, context deadline in future",
+			ctxTimeout: ptrutil.Ptr(10 * time.Second),
+			timeout:    0 * time.Second,
+			want:       true,
+		},
+		{
+			name:       "Timeout is negative, context deadline in future",
+			ctxTimeout: ptrutil.Ptr(10 * time.Second),
+			timeout:    -1 * time.Second,
+			want:       true,
+		},
+		{
+			name:       "Timeout provided, context has no deadline",
+			ctxTimeout: nil,
+			timeout:    1 * time.Second,
+			want:       true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			ctx := context.Background()
+			if test.ctxTimeout != nil {
+				var cancel context.CancelFunc
+
+				ctx, cancel = context.WithTimeout(ctx, *test.ctxTimeout)
+				defer cancel()
+			}
+
+			got := TimeoutWithinContext(ctx, test.timeout)
+			assert.Equal(t, test.want, got)
+		})
+	}
+}