Fix aggregate maxTimeMS integration test.

GODRIVER-813

Change-Id: Icdaba2d0768109ca17aba7f4dca1c69d6fc6606c

Author: Divjot Arora

internal/testutil/config.go

@@ -9,8 +9,10 @@ package testutil
 import (
 	"context"
 	"fmt"
+	"math"
 	"os"
 	"reflect"
+	"strconv"
 	"strings"
 	"sync"
 	"testing"
@@ -283,3 +285,29 @@ func Integration(t *testing.T) {
 		t.Skip("skipping integration test in short mode")
 	}
 }
+
+// compareVersions compares two version number strings (i.e. positive integers separated by
+// periods). Comparisons are done to the lesser precision of the two versions. For example, 3.2 is
+// considered equal to 3.2.11, whereas 3.2.0 is considered less than 3.2.11.
+//
+// Returns a positive int if version1 is greater than version2, a negative int if version1 is less
+// than version2, and 0 if version1 is equal to version2.
+func CompareVersions(t *testing.T, v1 string, v2 string) int {
+	n1 := strings.Split(v1, ".")
+	n2 := strings.Split(v2, ".")
+
+	for i := 0; i < int(math.Min(float64(len(n1)), float64(len(n2)))); i++ {
+		i1, err := strconv.Atoi(n1[i])
+		require.NoError(t, err)
+
+		i2, err := strconv.Atoi(n2[i])
+		require.NoError(t, err)
+
+		difference := i1 - i2
+		if difference != 0 {
+			return difference
+		}
+	}
+
+	return 0
+}

x/mongo/driver/integration/aggregate_test.go

@@ -0,0 +1,151 @@
+// Copyright (C) MongoDB, Inc. 2017-present.
+//
+// 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
+
+package integration
+
+import (
+	"context"
+	"github.com/mongodb/mongo-go-driver/bson"
+	"github.com/mongodb/mongo-go-driver/event"
+	"github.com/mongodb/mongo-go-driver/internal/testutil"
+	"github.com/mongodb/mongo-go-driver/mongo/options"
+	"github.com/mongodb/mongo-go-driver/mongo/readpref"
+	"github.com/mongodb/mongo-go-driver/x/bsonx"
+	"github.com/mongodb/mongo-go-driver/x/mongo/driver"
+	"github.com/mongodb/mongo-go-driver/x/mongo/driver/session"
+	"github.com/mongodb/mongo-go-driver/x/mongo/driver/topology"
+	"github.com/mongodb/mongo-go-driver/x/mongo/driver/uuid"
+	"github.com/mongodb/mongo-go-driver/x/network/command"
+	"github.com/mongodb/mongo-go-driver/x/network/description"
+	"github.com/stretchr/testify/require"
+	"testing"
+	"time"
+)
+
+func setUpMonitor() (*event.CommandMonitor, chan *event.CommandStartedEvent, chan *event.CommandSucceededEvent, chan *event.CommandFailedEvent) {
+	started := make(chan *event.CommandStartedEvent, 1)
+	succeeded := make(chan *event.CommandSucceededEvent, 1)
+	failed := make(chan *event.CommandFailedEvent, 1)
+
+	return &event.CommandMonitor{
+		Started: func(ctx context.Context, e *event.CommandStartedEvent) {
+			started <- e
+		},
+		Succeeded: func(ctx context.Context, e *event.CommandSucceededEvent) {
+			succeeded <- e
+		},
+		Failed: func(ctx context.Context, e *event.CommandFailedEvent) {
+			failed <- e
+		},
+	}, started, succeeded, failed
+}
+
+func skipIfBelow32(ctx context.Context, t *testing.T, topo *topology.Topology) {
+	server, err := topo.SelectServer(ctx, description.WriteSelector())
+	noerr(t, err)
+
+	versionCmd := bsonx.Doc{{"serverStatus", bsonx.Int32(1)}}
+	serverStatus, err := testutil.RunCommand(t, server.Server, dbName, versionCmd)
+	version, err := serverStatus.LookupErr("version")
+
+	if testutil.CompareVersions(t, version.StringValue(), "3.2") < 0 {
+		t.Skip()
+	}
+}
+
+func TestAggregate(t *testing.T) {
+	t.Run("TestMaxTimeMSInGetMore", func(t *testing.T) {
+		ctx := context.Background()
+		monitor, started, succeeded, failed := setUpMonitor()
+		dbName := "TestAggMaxTimeDB"
+		collName := "TestAggMaxTimeColl"
+		top := testutil.MonitoredTopology(t, dbName, monitor)
+		clearChannels(started, succeeded, failed)
+		skipIfBelow32(ctx, t, top)
+
+		clientID, err := uuid.New()
+		noerr(t, err)
+
+		ns := command.Namespace{
+			DB:         dbName,
+			Collection: collName,
+		}
+		pool := &session.Pool{}
+
+		clearChannels(started, succeeded, failed)
+		_, err = driver.Insert(
+			ctx,
+			command.Insert{
+				NS: ns,
+				Docs: []bsonx.Doc{
+					{{"x", bsonx.Int32(1)}},
+					{{"x", bsonx.Int32(1)}},
+					{{"x", bsonx.Int32(1)}},
+				},
+			},
+			top,
+			description.WriteSelector(),
+			clientID,
+			pool,
+			false,
+		)
+		noerr(t, err)
+
+		clearChannels(started, succeeded, failed)
+		cmd := command.Aggregate{
+			NS:       ns,
+			Pipeline: bsonx.Arr{},
+		}
+		batchCursor, err := driver.Aggregate(
+			ctx,
+			cmd,
+			top,
+			description.ReadPrefSelector(readpref.Primary()),
+			description.WriteSelector(),
+			clientID,
+			pool,
+			bson.DefaultRegistry,
+			options.Aggregate().SetMaxAwaitTime(10*time.Millisecond).SetBatchSize(2),
+		)
+		noerr(t, err)
+
+		var e *event.CommandStartedEvent
+		select {
+		case e = <-started:
+		case <-time.After(200 * time.Millisecond):
+			t.Fatal("timed out waiting for aggregate")
+		}
+
+		require.Equal(t, "aggregate", e.CommandName)
+
+		clearChannels(started, succeeded, failed)
+		// first Next() should automatically return true
+		require.True(t, batchCursor.Next(ctx), "expected true from first Next, got false")
+		clearChannels(started, succeeded, failed)
+		batchCursor.Next(ctx) // should do getMore
+
+		select {
+		case e = <-started:
+		case <-time.After(200 * time.Millisecond):
+			t.Fatal("timed out waiting for getMore")
+		}
+		require.Equal(t, "getMore", e.CommandName)
+		_, err = e.Command.LookupErr("maxTimeMS")
+		noerr(t, err)
+	})
+}
+
+func clearChannels(s chan *event.CommandStartedEvent, succ chan *event.CommandSucceededEvent, f chan *event.CommandFailedEvent) {
+	for len(s) > 0 {
+		<-s
+	}
+	for len(succ) > 0 {
+		<-succ
+	}
+	for len(f) > 0 {
+		<-f
+	}
+}

x/network/integration/aggregate_test.go

@@ -8,24 +8,17 @@ package integration
 
 import (
 	"context"
-	"fmt"
-	"os"
 	"strings"
 	"testing"
-	"time"
 
 	"github.com/mongodb/mongo-go-driver/bson"
 	"github.com/mongodb/mongo-go-driver/internal/testutil"
-	"github.com/mongodb/mongo-go-driver/internal/testutil/israce"
 	"github.com/mongodb/mongo-go-driver/mongo/writeconcern"
 	"github.com/mongodb/mongo-go-driver/x/bsonx"
-	"github.com/mongodb/mongo-go-driver/x/bsonx/bsoncore"
-	"github.com/mongodb/mongo-go-driver/x/mongo/driver"
 	"github.com/mongodb/mongo-go-driver/x/mongo/driver/topology"
 	"github.com/mongodb/mongo-go-driver/x/network/address"
 	"github.com/mongodb/mongo-go-driver/x/network/command"
 	"github.com/mongodb/mongo-go-driver/x/network/description"
-	"github.com/stretchr/testify/assert"
 )
 
 func TestCommandAggregate(t *testing.T) {
@@ -158,146 +151,3 @@ func TestCommandAggregate(t *testing.T) {
 		noerr(t, err)
 	})
 }
-
-func TestAggregatePassesMaxAwaitTimeMSThroughToGetMore(t *testing.T) {
-	if os.Getenv("TOPOLOGY") != "replica_set" {
-		t.Skip()
-	}
-
-	startedChan, succeededChan, failedChan, monitor := initMonitor()
-
-	dbName := fmt.Sprintf("mongo-go-driver-%d-agg", os.Getpid())
-	colName := testutil.ColName(t)
-
-	server, err := testutil.MonitoredTopology(t, dbName, monitor).SelectServer(context.Background(), description.WriteSelector())
-	noerr(t, err)
-
-	versionCmd := bsonx.Doc{{"serverStatus", bsonx.Int32(1)}}
-	serverStatus, err := testutil.RunCommand(t, server.Server, dbName, versionCmd)
-	version, err := serverStatus.LookupErr("version")
-
-	if compareVersions(t, version.StringValue(), "3.6") < 0 {
-		t.Skip()
-	}
-
-	// create capped collection
-	createCmd := bsonx.Doc{
-		{"create", bsonx.String(colName)},
-		{"capped", bsonx.Boolean(true)},
-		{"size", bsonx.Int32(1000)}}
-	_, err = testutil.RunCommand(t, server.Server, dbName, createCmd)
-	noerr(t, err)
-
-	conn, err := server.Connection(context.Background())
-	noerr(t, err)
-
-	// create an aggregate command that results with a TAILABLEAWAIT cursor
-	result, err := (&command.Aggregate{
-		NS: command.Namespace{DB: dbName, Collection: testutil.ColName(t)},
-		Pipeline: bsonx.Arr{
-			bsonx.Document(bsonx.Doc{
-				{"$changeStream", bsonx.Document(bsonx.Doc{})}}),
-			bsonx.Document(bsonx.Doc{
-				{"$match", bsonx.Document(bsonx.Doc{
-					{"fullDocument._id", bsonx.Document(bsonx.Doc{{"$gte", bsonx.Int32(1)}})},
-				})}})},
-		Opts: []bsonx.Elem{{"batchSize", bsonx.Int32(2)}},
-		CursorOpts: []bsonx.Elem{
-			{"batchSize", bsonx.Int32(2)},
-			{"maxTimeMS", bsonx.Int64(50)},
-		},
-	}).RoundTrip(context.Background(), server.SelectedDescription(), conn)
-	noerr(t, err)
-
-	cursor, err := driver.NewBatchCursor(
-		bsoncore.Document(result), nil, nil, server.Server,
-		bsonx.Elem{"batchSize", bsonx.Int32(2)}, bsonx.Elem{"maxTimeMS", bsonx.Int64(50)},
-	)
-	noerr(t, err)
-
-	// insert some documents
-	insertCmd := bsonx.Doc{
-		{"insert", bsonx.String(colName)},
-		{"documents", bsonx.Array(bsonx.Arr{
-			bsonx.Document(bsonx.Doc{{"_id", bsonx.Int32(1)}}),
-			bsonx.Document(bsonx.Doc{{"_id", bsonx.Int32(2)}}),
-			bsonx.Document(bsonx.Doc{{"_id", bsonx.Int32(3)}})})}}
-	_, err = testutil.RunCommand(t, server.Server, dbName, insertCmd)
-
-	// wait a bit between insert and getMore commands
-	time.Sleep(time.Millisecond * 100)
-	if israce.Enabled {
-		time.Sleep(time.Millisecond * 400) // wait a little longer when race detector is enabled.
-	}
-
-	ctx, cancel := context.WithCancel(context.Background())
-	if israce.Enabled {
-		time.AfterFunc(time.Millisecond*2000, cancel)
-	} else {
-		time.AfterFunc(time.Millisecond*900, cancel)
-	}
-	for cursor.Next(ctx) {
-	}
-
-	// allow for iteration over range chan
-	close(startedChan)
-	close(succeededChan)
-	close(failedChan)
-
-	// no commands should have failed
-	if len(failedChan) != 0 {
-		t.Errorf("%d command(s) failed", len(failedChan))
-	}
-
-	// check the expected commands were started
-	for started := range startedChan {
-		switch started.CommandName {
-		case "aggregate":
-			assert.Equal(t, 2, int(started.Command.Lookup("cursor", "batchSize").Int32()))
-			assert.Equal(t, started.Command.Lookup("maxAwaitTimeMS"), bson.RawValue{},
-				"Should not have sent maxAwaitTimeMS in find command")
-		case "getMore":
-			assert.Equal(t, 2, int(started.Command.Lookup("batchSize").Int32()))
-			assert.Equal(t, 50, int(started.Command.Lookup("maxTimeMS").Int64()),
-				"Should have sent maxTimeMS in getMore command")
-		default:
-			continue
-		}
-	}
-
-	// to keep track of seen documents
-	id := 1
-
-	// check expected commands succeeded
-	for succeeded := range succeededChan {
-		switch succeeded.CommandName {
-		case "aggregate":
-			assert.Equal(t, 1, int(succeeded.Reply.Lookup("ok").Double()))
-
-			actual, err := succeeded.Reply.Lookup("cursor", "firstBatch").Array().Values()
-			assert.NoError(t, err)
-
-			for _, v := range actual {
-				assert.Equal(t, id, int(v.Document().Lookup("fullDocument", "_id").Int32()))
-				id++
-			}
-		case "getMore":
-			assert.Equal(t, "getMore", succeeded.CommandName)
-			assert.Equal(t, 1, int(succeeded.Reply.Lookup("ok").Double()))
-
-			actual, err := succeeded.Reply.Lookup("cursor", "nextBatch").Array().Values()
-			assert.NoError(t, err)
-
-			for _, v := range actual {
-				assert.Equal(t, id, int(v.Document().Lookup("fullDocument", "_id").Int32()))
-				id++
-			}
-		default:
-			continue
-		}
-	}
-
-	if id <= 3 {
-		t.Errorf("not all documents returned; last seen id = %d", id-1)
-	}
-}