garage-k2v-go/read_batch_test.go

package k2v_test

import (
	k2v "code.notaphish.fyi/milas/garage-k2v-go"
	"fmt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"math/rand/v2"
	"strconv"
	"strings"
	"testing"
)

func TestClient_ReadBatch(t *testing.T) {
	f, ctx := newFixture(t)

	pk1 := randomKey("pk1")
	sk1 := randomKey("sk1")
	require.NoError(t, f.cli.InsertItem(ctx, f.bucket, pk1, sk1, "", []byte(strings.Join([]string{"hello", pk1, sk1}, "-"))))

	pk2 := randomKey("pk2")
	sk2 := randomKey("sk2")
	for i := range 5 {
		require.NoError(t, f.cli.InsertItem(ctx, f.bucket, pk2, sk2, "", []byte(strings.Join([]string{"hello", pk2, sk2, strconv.Itoa(i)}, "-"))))
	}

	pk3 := randomKey("pk3")
	for i := range 5 {
		skN := randomKey(fmt.Sprintf("sk%d", i+3))
		require.NoError(t, f.cli.InsertItem(ctx, f.bucket, pk3, skN, "", []byte(strings.Join([]string{"hello", pk3, skN, strconv.Itoa(i)}, "-"))))
	}

	q := []k2v.ReadBatchSearch{
		{
			PartitionKey: pk1,
		},
		{
			PartitionKey: pk2,
			SingleItem:   true,
			Start:        sk2,
		},
		{
			PartitionKey: pk3,
		},
	}

	results, err := f.cli.ReadBatch(ctx, f.bucket, q)
	require.NoError(t, err)
	require.NotEmpty(t, results)
	require.Len(t, results, 3)

	assert.Equal(t, pk1, results[0].PartitionKey)
	if assert.Len(t, results[0].Items, 1) && assert.Len(t, results[0].Items[0].Values, 1) {
		assert.Equal(t, sk1, results[0].Items[0].SortKey)
		assert.NotEmpty(t, results[0].Items[0].CausalityToken)
		assert.Contains(t, results[0].Items[0].Values[0].GoString(), "hello")
	}

	assert.Equal(t, pk2, results[1].PartitionKey)
	if assert.Len(t, results[1].Items, 1) && assert.Len(t, results[1].Items[0].Values, 5) {
		assert.Equal(t, sk2, results[1].Items[0].SortKey)
		assert.NotEmpty(t, results[1].Items[0].CausalityToken)
		for i := range results[1].Items[0].Values {
			assert.Contains(t, results[1].Items[0].Values[i].GoString(), "hello")
		}
	}

	assert.Equal(t, pk3, results[2].PartitionKey)
	if assert.Len(t, results[2].Items, 5) {
		for _, item := range results[2].Items {
			assert.NotEmpty(t, item.SortKey)
			assert.NotEmpty(t, item.CausalityToken)
			if assert.Len(t, item.Values, 1) {
				assert.Contains(t, item.Values[0].GoString(), "hello")
			}
		}
	}
}

func TestBulkGet(t *testing.T) {
	f, ctx := newFixture(t)

	keys := make([]k2v.ItemKey, 500)
	for i := range keys {
		keys[i] = k2v.ItemKey{
			PartitionKey: randomPk(),
			SortKey:      randomSk(),
		}
		require.NoError(t, f.cli.InsertItem(ctx, f.bucket, keys[i].PartitionKey, keys[i].SortKey, "", []byte("hello"+strconv.Itoa(i))))
	}

	rand.Shuffle(len(keys), func(i, j int) {
		keys[i], keys[j] = keys[j], keys[i]
	})

	items, err := k2v.BulkGet(ctx, f.cli, f.bucket, keys)
	require.NoError(t, err)
	require.NotEmpty(t, items)

	require.Equal(t, len(keys), len(items))
	for i := range keys {
		require.Equal(t, keys[i].SortKey, items[i].SortKey)
		require.Len(t, items[i].Values, 1)
		require.Contains(t, string(items[i].Values[0]), "hello")
	}
}