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