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

package objectstore_test

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

	"queryorchestration/internal/serviceconfig"
	objectstore "queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/test"
	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) {
	os.Clearenv()
	ctx := context.Background()
	c := objectstore.ObjectStoreConfig{}
	err := c.SetStoreClient(ctx)
	require.NoError(t, err)
	assert.NotNil(t, c.StoreClient)
}

func TestPingStoreByName(t *testing.T) {
	ctx := context.Background()
	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(context.Background(), 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 TestCalculateETag(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping long test in short mode")
	}
	ctx := context.Background()

	cfg := &StoreConfig{}
	test.SetCfgProvider(t, cfg)
	acfg, clean := test.CreateAWSContainer(t, ctx, cfg, &test.CreateAWSConfig{})
	defer clean()
	test.SetStoreClient(t, ctx, cfg, acfg.ExternalEndpoint)
	test.CreateBucket(t, ctx, cfg, test.BucketName)

	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.BucketName),
		Key:    aws.String("example"),
		Body:   uploadReader,
	})
	require.NoError(t, err)

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

func TestGetBucketKey(t *testing.T) {
	key := objectstore.BucketKey{}
	assert.Equal(t, "//00010101/", key.String())
	key.CreatedAt = time.Date(2025, time.March, 31, 0, 0, 0, 0, time.UTC)
	assert.Equal(t, "//20250331/", key.String())
	key.ClientID = "AAA"
	assert.Equal(t, "AAA//20250331/", key.String())
	key.Filename = "coolfile"
	assert.Equal(t, "AAA//20250331/coolfile", key.String())
	key.Location = objectstore.Import
	assert.Equal(t, "AAA/import/20250331/coolfile", key.String())
	part := uint8(3)
	key.Part = &part
	assert.Equal(t, "AAA/import/20250331/3/coolfile", key.String())
	key.Location = objectstore.TextTextract
	assert.Equal(t, "AAA/text/textract/20250331/3/coolfile", key.String())
	key.Location = objectstore.TextOut
	assert.Equal(t, "AAA/text/out/20250331/3/coolfile", key.String())
	key.Location = objectstore.Export
	assert.Equal(t, "AAA/export/20250331/3/coolfile", key.String())
}
func TestParseBucketKey(t *testing.T) {
	_, err := objectstore.ParseBucketKey("")
	assert.Error(t, err)

	_, err = objectstore.ParseBucketKey("AAA")
	assert.Error(t, err)

	_, err = objectstore.ParseBucketKey("//00010101/")
	assert.Error(t, err)

	_, err = objectstore.ParseBucketKey("//20250331/")
	assert.Error(t, err)

	_, err = objectstore.ParseBucketKey("AAA//20250331/")
	assert.Error(t, err)

	_, err = objectstore.ParseBucketKey("AAA//20250331/coolfile")
	assert.Error(t, err)

	key, err := objectstore.ParseBucketKey("AAA/import/20250331/coolfile")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.Import,
		Filename:  "coolfile",
		CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC),
	}, key)

	key, err = objectstore.ParseBucketKey("AAA/text/textract/20250331/coolfile")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.TextTextract,
		Filename:  "coolfile",
		CreatedAt: time.Date(2025, 2, 31, 0, 0, 0, 0, time.UTC),
	}, key)

	key, err = objectstore.ParseBucketKey("AAA/text/out/20250331/coolfile")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.TextOut,
		Filename:  "coolfile",
		CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC),
	}, key)

	key, err = objectstore.ParseBucketKey("AAA/export/20250331/coolfile")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.Export,
		Filename:  "coolfile",
		CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC),
	}, key)

	part := uint8(3)
	key, err = objectstore.ParseBucketKey("AAA/import/20250331/3/coolfile")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.Import,
		Filename:  "coolfile",
		CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC),
		Part:      &part,
	}, key)

	key, err = objectstore.ParseBucketKey("AAA/text/textract/20250331/3/coolfile")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.TextTextract,
		Filename:  "coolfile",
		CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC),
		Part:      &part,
	}, key)

	key, err = objectstore.ParseBucketKey("AAA/text/out/20250331/3/coolfile")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.TextOut,
		Filename:  "coolfile",
		CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC),
		Part:      &part,
	}, key)

	key, err = objectstore.ParseBucketKey("AAA/export/20250331/3/coolfile")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.Export,
		Filename:  "coolfile",
		CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC),
		Part:      &part,
	}, key)

	key, err = objectstore.ParseBucketKey("AAA/import/20250331/cool-file")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.Import,
		Filename:  "cool-file",
		CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC),
	}, key)

	key, err = objectstore.ParseBucketKey("AAA/import/20250331/cool_file")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.Import,
		Filename:  "cool_file",
		CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC),
	}, key)

	key, err = objectstore.ParseBucketKey("AAA/import/20250331/coolfile365")
	assert.NoError(t, err)
	assert.EqualExportedValues(t, objectstore.BucketKey{
		ClientID:  "AAA",
		Location:  objectstore.Import,
		Filename:  "coolfile365",
		CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC),
	}, key)

	_, err = objectstore.ParseBucketKey("AAA/import/99999999/coolfile")
	assert.Error(t, err)

	_, err = objectstore.ParseBucketKey("AAA/import/20250331/99999999/coolfile")
	assert.Error(t, err)

	_, err = objectstore.ParseBucketKey("AAA/import/20250331/999/coolfile")
	assert.Error(t, err)

	_, err = objectstore.ParseBucketKey("AAA/import/20250331/1/cool**file")
	assert.Error(t, err)
}