github.com/sequix/cortex@v1.1.6/pkg/chunk/storage/object_client_test.go (about)

     1  package storage
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"sort"
     7  	"strconv"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/sequix/cortex/pkg/chunk"
    14  	"github.com/sequix/cortex/pkg/chunk/testutils"
    15  	"github.com/prometheus/common/model"
    16  )
    17  
    18  func TestChunksBasic(t *testing.T) {
    19  	forAllFixtures(t, func(t *testing.T, _ chunk.IndexClient, client chunk.ObjectClient) {
    20  		const batchSize = 5
    21  		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
    22  		defer cancel()
    23  
    24  		// Write a few batches of chunks.
    25  		written := []string{}
    26  		for i := 0; i < 5; i++ {
    27  			keys, chunks, err := testutils.CreateChunks(i, batchSize, model.Now())
    28  			require.NoError(t, err)
    29  			written = append(written, keys...)
    30  			err = client.PutChunks(ctx, chunks)
    31  			require.NoError(t, err)
    32  		}
    33  
    34  		// Get a few batches of chunks.
    35  		for batch := 0; batch < 50; batch++ {
    36  			keysToGet := map[string]struct{}{}
    37  			chunksToGet := []chunk.Chunk{}
    38  			for len(chunksToGet) < batchSize {
    39  				key := written[rand.Intn(len(written))]
    40  				if _, ok := keysToGet[key]; ok {
    41  					continue
    42  				}
    43  				keysToGet[key] = struct{}{}
    44  				chunk, err := chunk.ParseExternalKey(userID, key)
    45  				require.NoError(t, err)
    46  				chunksToGet = append(chunksToGet, chunk)
    47  			}
    48  
    49  			chunksWeGot, err := client.GetChunks(ctx, chunksToGet)
    50  			require.NoError(t, err)
    51  			require.Equal(t, len(chunksToGet), len(chunksWeGot))
    52  
    53  			sort.Sort(ByKey(chunksToGet))
    54  			sort.Sort(ByKey(chunksWeGot))
    55  			for i := 0; i < len(chunksWeGot); i++ {
    56  				require.Equal(t, chunksToGet[i].ExternalKey(), chunksWeGot[i].ExternalKey(), strconv.Itoa(i))
    57  			}
    58  		}
    59  	})
    60  }