GODRIVER-956 fix legacyGetMore

Change-Id: I28a7702b5f00f6de99cab45e9780c43574cc5b58

Author: iwysiu

mongo/collection_internal_test.go

@@ -1672,6 +1672,38 @@ func TestCollection_Find_found(t *testing.T) {
 	require.Equal(t, results, []int{1, 2, 3, 4, 5})
 }
 
+func TestCollection_Find_With_Limit_And_BatchSize(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test in short mode")
+	}
+
+	coll := createTestCollection(t, nil, nil)
+	initCollection(t, coll)
+
+	batchSizes := []int32{2, 3, 4}
+	for _, b := range batchSizes {
+		name := fmt.Sprintf("Limit: 3, BatchSize: %v", b)
+		t.Run(name, func(t *testing.T) {
+			cursor, err := coll.Find(context.Background(), bson.D{}, options.Find().SetLimit(3).SetBatchSize(b))
+
+			numRecieved := 0
+			var doc bson.Raw
+			for cursor.Next(context.Background()) {
+				err = cursor.Decode(&doc)
+				require.NoError(t, err)
+
+				_, err = doc.LookupErr("_id")
+				require.NoError(t, err)
+
+				numRecieved++
+			}
+
+			require.NoError(t, cursor.Err())
+			require.Equal(t, 3, numRecieved)
+		})
+	}
+}
+
 func TestCollection_Find_notFound(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skipping integration test in short mode")

x/mongo/driver/batch_cursor.go

@@ -342,6 +342,13 @@ func (bc *BatchCursor) legacyGetMore(ctx context.Context) {
 	numToReturn := bc.batchSize
 	if bc.limit != 0 && bc.numReturned+bc.batchSize > bc.limit {
 		numToReturn = bc.limit - bc.numReturned
+		if numToReturn <= 0 {
+			err = bc.Close(ctx)
+			if err != nil {
+				bc.err = err
+			}
+			return
+		}
 	}
 	gm := wiremessage.GetMore{
 		FullCollectionName: bc.namespace.DB + "." + bc.namespace.Collection,