github.com/thanos-io/thanos@v0.32.5/pkg/cacheutil/cacheutil_test.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package cacheutil
     5  
     6  import (
     7  	"context"
     8  	"os"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/onsi/gomega"
    13  	"github.com/onsi/gomega/gleak"
    14  	"github.com/prometheus/client_golang/prometheus"
    15  	"go.uber.org/atomic"
    16  
    17  	"github.com/efficientgo/core/testutil"
    18  	"github.com/thanos-io/thanos/pkg/gate"
    19  )
    20  
    21  func TestMain(m *testing.M) {
    22  	g := gomega.NewGomega(func(message string, callerSkip ...int) {
    23  		panic(message)
    24  	})
    25  	code := m.Run()
    26  	g.Eventually(gleak.Goroutines).WithTimeout(time.Second * 20).ShouldNot(gleak.HaveLeaked())
    27  	os.Exit(code)
    28  }
    29  
    30  func TestDoWithBatch(t *testing.T) {
    31  	tests := map[string]struct {
    32  		items           []string
    33  		batchSize       int
    34  		expectedBatches int
    35  		concurrency     gate.Gate
    36  	}{
    37  		"no items": {
    38  			items:           []string{},
    39  			batchSize:       2,
    40  			expectedBatches: 0,
    41  			concurrency:     nil,
    42  		},
    43  
    44  		"fewer than batch size": {
    45  			items:           []string{"key1"},
    46  			batchSize:       2,
    47  			expectedBatches: 1,
    48  			concurrency:     nil,
    49  		},
    50  
    51  		"perfect sized for batch": {
    52  			items:           []string{"key1", "key2", "key3", "key4"},
    53  			batchSize:       2,
    54  			expectedBatches: 2,
    55  			concurrency:     nil,
    56  		},
    57  
    58  		"odd sized for batch": {
    59  			items:           []string{"key1", "key2", "key3", "key4", "key5"},
    60  			batchSize:       2,
    61  			expectedBatches: 3,
    62  			concurrency:     nil,
    63  		},
    64  
    65  		"odd sized with concurrency limit": {
    66  			items:           []string{"key1", "key2", "key3", "key4", "key5"},
    67  			batchSize:       2,
    68  			expectedBatches: 3,
    69  			concurrency:     gate.New(prometheus.NewPedanticRegistry(), 1, gate.Queries),
    70  		},
    71  	}
    72  
    73  	for testName, testData := range tests {
    74  		t.Run(testName, func(t *testing.T) {
    75  			actualBatches := atomic.Int64{}
    76  			_ = doWithBatch(context.Background(), len(testData.items), testData.batchSize, testData.concurrency, func(startIndex, endIndex int) error {
    77  				actualBatches.Inc()
    78  				return nil
    79  			})
    80  
    81  			testutil.Equals(t, int64(testData.expectedBatches), actualBatches.Load())
    82  		})
    83  	}
    84  }