add tests for odd number of sessions

Author: Divjot Arora

mongo/client_test.go

@@ -9,6 +9,7 @@ package mongo
 import (
 	"context"
 	"errors"
+	"math"
 	"testing"
 	"time"
 
@@ -261,70 +262,85 @@ func TestClient(t *testing.T) {
 	})
 	t.Run("endSessions", func(t *testing.T) {
 		cs := testutil.ConnString(t)
-
-		var started []*event.CommandStartedEvent
-		var failureReasons []string
-		cmdMonitor := &event.CommandMonitor{
-			Started: func(_ context.Context, evt *event.CommandStartedEvent) {
-				if evt.CommandName == "endSessions" {
-					started = append(started, evt)
-				}
-			},
-			Failed: func(_ context.Context, evt *event.CommandFailedEvent) {
-				if evt.CommandName == "endSessions" {
-					failureReasons = append(failureReasons, evt.Failure)
-				}
-			},
-		}
-		clientOpts := options.Client().ApplyURI(cs.Original).SetReadPreference(readpref.Primary()).
-			SetWriteConcern(writeconcern.New(writeconcern.WMajority())).SetMonitor(cmdMonitor)
-		client, err := Connect(bgCtx, clientOpts)
-		assert.Nil(t, err, "Connect error: %v", err)
-		defer func() {
-			_ = client.Disconnect(bgCtx)
-		}()
-
-		serverVersion, err := getServerVersion(client.Database("admin"))
-		assert.Nil(t, err, "getServerVersion error: %v", err)
-		if compareVersions(t, serverVersion, "3.6.0") < 1 {
-			t.Skip("skipping server version < 3.6")
-		}
-
-		coll := client.Database("foo").Collection("bar")
-		// Lower the batch size to force multiple batches.
 		originalBatchSize := endSessionsBatchSize
 		endSessionsBatchSize = 2
 		defer func() {
 			endSessionsBatchSize = originalBatchSize
-			_ = coll.Drop(bgCtx)
 		}()
 
-		// Do an application operation and create four sessions so endSessions will execute in two batches.
-		_, err = coll.CountDocuments(bgCtx, bson.D{})
-		assert.Nil(t, err, "CountDocuments error: %v", err)
-		numSessions := 4
-		var sessions []Session
-		for i := 0; i < numSessions; i++ {
-			sess, err := client.StartSession()
-			assert.Nil(t, err, "StartSession error at index %d: %v", i, err)
-			sessions = append(sessions, sess)
-		}
-		for _, sess := range sessions {
-			sess.EndSession(bgCtx)
+		testCases := []struct {
+			name            string
+			numSessions     int
+			eventBatchSizes []int
+		}{
+			{"number of sessions divides evenly", endSessionsBatchSize * 2, []int{endSessionsBatchSize, endSessionsBatchSize}},
+			{"number of sessions does not divide evenly", endSessionsBatchSize + 1, []int{endSessionsBatchSize, 1}},
 		}
+		for _, tc := range testCases {
+			t.Run(tc.name, func(t *testing.T) {
+				// Setup a client and skip the test based on server version.
+				var started []*event.CommandStartedEvent
+				var failureReasons []string
+				cmdMonitor := &event.CommandMonitor{
+					Started: func(_ context.Context, evt *event.CommandStartedEvent) {
+						if evt.CommandName == "endSessions" {
+							started = append(started, evt)
+						}
+					},
+					Failed: func(_ context.Context, evt *event.CommandFailedEvent) {
+						if evt.CommandName == "endSessions" {
+							failureReasons = append(failureReasons, evt.Failure)
+						}
+					},
+				}
+				clientOpts := options.Client().ApplyURI(cs.Original).SetReadPreference(readpref.Primary()).
+					SetWriteConcern(writeconcern.New(writeconcern.WMajority())).SetMonitor(cmdMonitor)
+				client, err := Connect(bgCtx, clientOpts)
+				assert.Nil(t, err, "Connect error: %v", err)
+				defer func() {
+					_ = client.Disconnect(bgCtx)
+				}()
+
+				serverVersion, err := getServerVersion(client.Database("admin"))
+				assert.Nil(t, err, "getServerVersion error: %v", err)
+				if compareVersions(t, serverVersion, "3.6.0") < 1 {
+					t.Skip("skipping server version < 3.6")
+				}
 
-		client.endSessions(bgCtx)
-		numEventsExpected := numSessions / endSessionsBatchSize
-		assert.Equal(t, len(started), numEventsExpected, "expected %d started events, got %d", numEventsExpected,
-			len(started))
-		assert.Equal(t, len(failureReasons), 0, "endSessions errors: %v", failureReasons)
+				coll := client.Database("foo").Collection("bar")
+				defer func() {
+					_ = coll.Drop(bgCtx)
+				}()
 
-		for i := 0; i < numEventsExpected; i++ {
-			sentArray := started[i].Command.Lookup("endSessions").Array()
-			values, _ := sentArray.Values()
-			assert.Equal(t, len(values), endSessionsBatchSize,
-				"batch size mismatch at index %d; expected %d sessions in batch, got %d", i, endSessionsBatchSize,
-				len(values))
+				// Do an application operation and create four sessions so endSessions will execute in two batches.
+				_, err = coll.CountDocuments(bgCtx, bson.D{})
+				assert.Nil(t, err, "CountDocuments error: %v", err)
+				var sessions []Session
+				for i := 0; i < tc.numSessions; i++ {
+					sess, err := client.StartSession()
+					assert.Nil(t, err, "StartSession error at index %d: %v", i, err)
+					sessions = append(sessions, sess)
+				}
+				for _, sess := range sessions {
+					sess.EndSession(bgCtx)
+				}
+
+				client.endSessions(bgCtx)
+				divisionResult := float64(tc.numSessions) / float64(endSessionsBatchSize)
+				numEventsExpected := int(math.Ceil(divisionResult))
+				assert.Equal(t, len(started), numEventsExpected, "expected %d started events, got %d", numEventsExpected,
+					len(started))
+				assert.Equal(t, len(failureReasons), 0, "endSessions errors: %v", failureReasons)
+
+				for i := 0; i < numEventsExpected; i++ {
+					sentArray := started[i].Command.Lookup("endSessions").Array()
+					values, _ := sentArray.Values()
+					expectedNumValues := tc.eventBatchSizes[i]
+					assert.Equal(t, len(values), expectedNumValues,
+						"batch size mismatch at index %d; expected %d sessions in batch, got %d", i, expectedNumValues,
+						len(values))
+				}
+			})
 		}
 	})
 }