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) }