batch_cursor_test.go

Documentation: go.mongodb.org/mongo-driver/x/mongo/driver

     1  // Copyright (C) MongoDB, Inc. 2022-present.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License"); you may
     4  // not use this file except in compliance with the License. You may obtain
     5  // a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
     6  
     7  package driver
     8  
     9  import (
    10  	"testing"
    11  	"time"
    12  
    13  	"go.mongodb.org/mongo-driver/internal/assert"
    14  )
    15  
    16  func TestBatchCursor(t *testing.T) {
    17  	t.Parallel()
    18  
    19  	t.Run("setBatchSize", func(t *testing.T) {
    20  		t.Parallel()
    21  
    22  		var size int32
    23  		bc := &BatchCursor{
    24  			batchSize: size,
    25  		}
    26  		assert.Equal(t, size, bc.batchSize, "expected batchSize %v, got %v", size, bc.batchSize)
    27  
    28  		size = int32(4)
    29  		bc.SetBatchSize(size)
    30  		assert.Equal(t, size, bc.batchSize, "expected batchSize %v, got %v", size, bc.batchSize)
    31  	})
    32  
    33  	t.Run("calcGetMoreBatchSize", func(t *testing.T) {
    34  		t.Parallel()
    35  
    36  		for _, tcase := range []struct {
    37  			name                               string
    38  			size, limit, numReturned, expected int32
    39  			ok                                 bool
    40  		}{
    41  			{
    42  				name:     "empty",
    43  				expected: 0,
    44  				ok:       true,
    45  			},
    46  			{
    47  				name:     "batchSize NEQ 0",
    48  				size:     4,
    49  				expected: 4,
    50  				ok:       true,
    51  			},
    52  			{
    53  				name:     "limit NEQ 0",
    54  				limit:    4,
    55  				expected: 0,
    56  				ok:       true,
    57  			},
    58  			{
    59  				name:        "limit NEQ and batchSize + numReturned EQ limit",
    60  				size:        4,
    61  				limit:       8,
    62  				numReturned: 4,
    63  				expected:    4,
    64  				ok:          true,
    65  			},
    66  			{
    67  				name:        "limit makes batchSize negative",
    68  				numReturned: 4,
    69  				limit:       2,
    70  				expected:    -2,
    71  				ok:          false,
    72  			},
    73  		} {
    74  			tcase := tcase
    75  			t.Run(tcase.name, func(t *testing.T) {
    76  				t.Parallel()
    77  
    78  				bc := &BatchCursor{
    79  					limit:       tcase.limit,
    80  					batchSize:   tcase.size,
    81  					numReturned: tcase.numReturned,
    82  				}
    83  
    84  				bc.SetBatchSize(tcase.size)
    85  
    86  				size, ok := calcGetMoreBatchSize(*bc)
    87  
    88  				assert.Equal(t, tcase.expected, size, "expected batchSize %v, got %v", tcase.expected, size)
    89  				assert.Equal(t, tcase.ok, ok, "expected ok %v, got %v", tcase.ok, ok)
    90  			})
    91  		}
    92  	})
    93  }
    94  
    95  func TestBatchCursorSetMaxTime(t *testing.T) {
    96  	t.Parallel()
    97  
    98  	tests := []struct {
    99  		name string
   100  		dur  time.Duration
   101  		want int64
   102  	}{
   103  		{
   104  			name: "empty",
   105  			dur:  0,
   106  			want: 0,
   107  		},
   108  		{
   109  			name: "partial milliseconds are truncated",
   110  			dur:  10_900 * time.Microsecond,
   111  			want: 10,
   112  		},
   113  		{
   114  			name: "millisecond input",
   115  			dur:  10 * time.Millisecond,
   116  			want: 10,
   117  		},
   118  	}
   119  
   120  	for _, test := range tests {
   121  		test := test
   122  
   123  		t.Run(test.name, func(t *testing.T) {
   124  			t.Parallel()
   125  
   126  			bc := BatchCursor{}
   127  			bc.SetMaxTime(test.dur)
   128  
   129  			got := bc.maxTimeMS
   130  			assert.Equal(t, test.want, got, "expected and actual maxTimeMS are different")
   131  		})
   132  	}
   133  }
   134
View as plain text