batches_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  
    12  	"github.com/google/go-cmp/cmp"
    13  	"go.mongodb.org/mongo-driver/internal/assert"
    14  	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
    15  )
    16  
    17  func TestBatches(t *testing.T) {
    18  	t.Run("Valid", func(t *testing.T) {
    19  		testCases := []struct {
    20  			name    string
    21  			batches *Batches
    22  			want    bool
    23  		}{
    24  			{"nil", nil, false},
    25  			{"missing identifier", &Batches{}, false},
    26  			{"no documents", &Batches{Identifier: "documents"}, false},
    27  			{"valid", &Batches{Identifier: "documents", Documents: make([]bsoncore.Document, 5)}, true},
    28  		}
    29  
    30  		for _, tc := range testCases {
    31  			t.Run(tc.name, func(t *testing.T) {
    32  				want := tc.want
    33  				got := tc.batches.Valid()
    34  				if got != want {
    35  					t.Errorf("Did not get expected result from Valid. got %t; want %t", got, want)
    36  				}
    37  			})
    38  		}
    39  	})
    40  	t.Run("ClearBatch", func(t *testing.T) {
    41  		batches := &Batches{Identifier: "documents", Current: make([]bsoncore.Document, 2, 10)}
    42  		if len(batches.Current) != 2 {
    43  			t.Fatalf("Length of current batch should be 2, but is %d", len(batches.Current))
    44  		}
    45  		batches.ClearBatch()
    46  		if len(batches.Current) != 0 {
    47  			t.Fatalf("Length of current batch should be 0, but is %d", len(batches.Current))
    48  		}
    49  	})
    50  	t.Run("AdvanceBatch", func(t *testing.T) {
    51  		documents := make([]bsoncore.Document, 0)
    52  		for i := 0; i < 5; i++ {
    53  			doc := make(bsoncore.Document, 100)
    54  			documents = append(documents, doc)
    55  		}
    56  
    57  		testCases := []struct {
    58  			name            string
    59  			batches         *Batches
    60  			maxCount        int
    61  			targetBatchSize int
    62  			maxDocSize      int
    63  			err             error
    64  			want            *Batches
    65  		}{
    66  			{
    67  				"current batch non-zero",
    68  				&Batches{Current: make([]bsoncore.Document, 2, 10)},
    69  				0, 0, 0, nil,
    70  				&Batches{Current: make([]bsoncore.Document, 2, 10)},
    71  			},
    72  			{
    73  				// all of the documents in the batch fit in targetBatchSize so the batch is created successfully
    74  				"documents fit in targetBatchSize",
    75  				&Batches{Documents: documents},
    76  				10, 600, 1000, nil,
    77  				&Batches{Documents: documents[:0], Current: documents[0:]},
    78  			},
    79  			{
    80  				// the first doc is bigger than targetBatchSize but smaller than maxDocSize so it is taken alone
    81  				"first document larger than targetBatchSize, smaller than maxDocSize",
    82  				&Batches{Documents: documents},
    83  				10, 5, 100, nil,
    84  				&Batches{Documents: documents[1:], Current: documents[:1]},
    85  			},
    86  		}
    87  
    88  		for _, tc := range testCases {
    89  			t.Run(tc.name, func(t *testing.T) {
    90  				err := tc.batches.AdvanceBatch(tc.maxCount, tc.targetBatchSize, tc.maxDocSize)
    91  				if !cmp.Equal(err, tc.err, cmp.Comparer(compareErrors)) {
    92  					t.Errorf("Errors do not match. got %v; want %v", err, tc.err)
    93  				}
    94  				if !cmp.Equal(tc.batches, tc.want) {
    95  					t.Errorf("Batches is not in correct state after AdvanceBatch. got %v; want %v", tc.batches, tc.want)
    96  				}
    97  			})
    98  		}
    99  
   100  		t.Run("middle document larger than targetBatchSize, smaller than maxDocSize", func(t *testing.T) {
   101  			// a batch is made but one document is too big, so everything before it is taken.
   102  			// on the second call to AdvanceBatch, only the large document is taken
   103  
   104  			middleLargeDoc := make([]bsoncore.Document, 0)
   105  			for i := 0; i < 5; i++ {
   106  				doc := make(bsoncore.Document, 100)
   107  				middleLargeDoc = append(middleLargeDoc, doc)
   108  			}
   109  			largeDoc := make(bsoncore.Document, 900)
   110  			middleLargeDoc[2] = largeDoc
   111  			batches := &Batches{Documents: middleLargeDoc}
   112  			maxCount := 10
   113  			targetSize := 600
   114  			maxDocSize := 1000
   115  
   116  			// first batch should take first 2 docs (size 100 each)
   117  			err := batches.AdvanceBatch(maxCount, targetSize, maxDocSize)
   118  			assert.Nil(t, err, "AdvanceBatch error: %v", err)
   119  			want := &Batches{Current: middleLargeDoc[:2], Documents: middleLargeDoc[2:]}
   120  			assert.Equal(t, want, batches, "expected batches %v, got %v", want, batches)
   121  
   122  			// second batch should take single large doc (size 900)
   123  			batches.ClearBatch()
   124  			err = batches.AdvanceBatch(maxCount, targetSize, maxDocSize)
   125  			assert.Nil(t, err, "AdvanceBatch error: %v", err)
   126  			want = &Batches{Current: middleLargeDoc[2:3], Documents: middleLargeDoc[3:]}
   127  			assert.Equal(t, want, batches, "expected batches %v, got %v", want, batches)
   128  
   129  			// last batch should take last 2 docs (size 100 each)
   130  			batches.ClearBatch()
   131  			err = batches.AdvanceBatch(maxCount, targetSize, maxDocSize)
   132  			assert.Nil(t, err, "AdvanceBatch error: %v", err)
   133  			want = &Batches{Current: middleLargeDoc[3:], Documents: middleLargeDoc[:0]}
   134  			assert.Equal(t, want, batches, "expected batches %v, got %v", want, batches)
   135  		})
   136  	})
   137  }
   138
View as plain text