query-orchestration/internal/serviceconfig/objectstore/config_test.go

package objectstore_test

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/test"

	objectstore "queryorchestration/internal/serviceconfig/objectstore"

	objectstoremock "queryorchestration/mocks/objectstore"

	awsc "queryorchestration/internal/serviceconfig/aws"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

func TestGetStoreClient(t *testing.T) {
	c := objectstore.ObjectStoreConfig{}
	assert.Nil(t, c.GetStoreClient())
	c.StoreClient = &s3.Client{}
	assert.Equal(t, &s3.Client{}, c.GetStoreClient())
}

func TestStoreClient(t *testing.T) {
	ctx := t.Context()
	c := objectstore.ObjectStoreConfig{}
	err := c.SetStoreClient(ctx)
	require.NoError(t, err)
	assert.NotNil(t, c.StoreClient)
}

func TestPingStoreByName(t *testing.T) {
	ctx := t.Context()
	mocks3 := objectstoremock.NewMockS3Client(t)
	c := objectstore.ObjectStoreConfig{
		StoreClient: mocks3,
	}
	name := "bucket-name"

	mocks3.EXPECT().
		HeadBucket(
			mock.Anything,
			mock.MatchedBy(func(in *s3.HeadBucketInput) bool {
				return *in.Bucket == name
			}),
			mock.Anything,
		).
		Return(&s3.HeadBucketOutput{}, nil)

	err := c.PingStoreByName(ctx, name)
	require.NoError(t, err)
}

func TestGetS3Client(t *testing.T) {
	c := objectstore.ObjectStoreConfig{}
	assert.Empty(t, c.GetS3Endpoint())
	c.AWSEndpointUrlS3 = "s3_endpoint"
	assert.Equal(t, "s3_endpoint", c.GetS3Endpoint())
}

type StoreConfig struct {
	serviceconfig.BaseConfig
	objectstore.ObjectStoreConfig
	awsc.AWSConfig
}

func TestSetStoreClientAndPingByName(t *testing.T) {
	ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer cancel()
	cfg := &StoreConfig{}
	name := "bucket"

	err := cfg.SetStoreClientAndPingByName(ctx, name)

	if !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("Expected timeout error, got: %v", err)
	}
}

func TestCalculateAndCompareETag(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	t.SkipNow()
	ctx := t.Context()

	cfg := &StoreConfig{}

	acfg := test.CreateAWSContainer(t, cfg)
	test.SetStoreClient(t, ctx, cfg, acfg.ExternalEndpoint)
	test.CreateBucket(t, cfg)

	testContent := "hi"
	hashReader := strings.NewReader(testContent)
	uploadReader := strings.NewReader(testContent)

	hash, err := cfg.CalculateETag(ctx, hashReader)
	require.NoError(t, err)

	res, err := cfg.StoreClient.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(test.GetBucketName(t)),
		Key:    aws.String("example"),
		Body:   uploadReader,
	})
	require.NoError(t, err)

	assert.Equal(t, *res.ETag, hash)
}

func TestGetDirectoryPart(t *testing.T) {
	cfg := &StoreConfig{}

	assert.Equal(t, uint16(0), cfg.GetDirectoryPart(0, 0))
	assert.Equal(t, uint16(0), cfg.GetDirectoryPart(0, 1))
	assert.Equal(t, uint16(0), cfg.GetDirectoryPart(0, 2))
	assert.Equal(t, uint16(1), cfg.GetDirectoryPart(0, 1000))

	cfg.MaxDirSize = 2
	assert.Equal(t, uint16(0), cfg.GetDirectoryPart(0, 0))
	assert.Equal(t, uint16(0), cfg.GetDirectoryPart(0, 1))
	assert.Equal(t, uint16(1), cfg.GetDirectoryPart(0, 2))
	assert.Equal(t, uint16(1), cfg.GetDirectoryPart(1, 0))
	assert.Equal(t, uint16(1), cfg.GetDirectoryPart(1, 1))
	assert.Equal(t, uint16(2), cfg.GetDirectoryPart(1, 2))
	assert.Equal(t, uint16(2), cfg.GetDirectoryPart(2, 0))
	assert.Equal(t, uint16(2), cfg.GetDirectoryPart(2, 1))
	assert.Equal(t, uint16(3), cfg.GetDirectoryPart(2, 2))
	assert.Equal(t, uint16(3), cfg.GetDirectoryPart(3, 0))
	assert.Equal(t, uint16(3), cfg.GetDirectoryPart(3, 1))
	assert.Equal(t, uint16(4), cfg.GetDirectoryPart(3, 2))
}

func TestSetAWSEndpointUrlS3(t *testing.T) {
	cfg := objectstore.ObjectStoreConfig{}

	assert.Equal(t, "", cfg.AWSEndpointUrlS3)
	newVal := "val"
	cfg.SetS3Endpoint(newVal)
	assert.Equal(t, "val", cfg.AWSEndpointUrlS3)
}

func TestSetAWSSessionToken(t *testing.T) {
	cfg := objectstore.ObjectStoreConfig{}

	assert.False(t, cfg.AWSS3UsePathStyle)
	newVal := true
	cfg.SetS3UsePathStyle(newVal)
	assert.True(t, cfg.AWSS3UsePathStyle)
}