Files
query-orchestration/internal/test/objectstore.go
T
Michael McGuinness cc2278086f Merged in feature/import (pull request #157)
Feature/import

* importstart

* pp

* tests

* importtests

* 100GB

* lint

* passtests

* utc

* awsprofile

* doublequotes

* host
2025-05-27 15:28:46 +00:00

93 lines
2.6 KiB
Go

package test
import (
"context"
"io"
"log/slog"
"testing"
objectstore "queryorchestration/internal/serviceconfig/objectstore"
awsc "queryorchestration/internal/serviceconfig/aws"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/s3"
awstypes "github.com/aws/aws-sdk-go-v2/service/s3/types"
"github.com/stretchr/testify/require"
)
func GetBucketName(t testing.TB) string {
return GetAlias(t, "bucket")
}
func CreateBucket(t testing.TB, cfg objectstore.ConfigProvider) {
t.Helper()
bucketName := GetBucketName(t)
_, err := cfg.GetStoreClient().CreateBucket(t.Context(), &s3.CreateBucketInput{
Bucket: aws.String(bucketName),
})
require.NoError(t, err)
slog.Info("create bucket", "bucket", bucketName)
err = cfg.PingStoreByName(t.Context(), bucketName)
require.NoError(t, err)
cfg.SetBucket(bucketName)
}
func SetBucketNotifs(t testing.TB, ctx context.Context, cfg objectstore.ConfigProvider) {
t.Helper()
bucketName := GetBucketName(t)
arn := GetQueueArn(t, cfg, StoreEventRunnerName)
_, err := cfg.GetStoreClient().PutBucketNotificationConfiguration(ctx, &s3.PutBucketNotificationConfigurationInput{
Bucket: &bucketName,
NotificationConfiguration: &awstypes.NotificationConfiguration{
QueueConfigurations: []awstypes.QueueConfiguration{
{
QueueArn: &arn,
Events: []awstypes.Event{
awstypes.EventS3ObjectCreated,
},
},
},
},
})
require.NoError(t, err)
slog.Info("set bucket notifs", "bucket", bucketName, "queueARN", arn)
}
func SetStoreClient(t testing.TB, ctx context.Context, cfg objectstore.ConfigProvider, endpoint string) {
t.Helper()
acfg, err := awsc.GetAWSConfigWithOpts(ctx, func(lo *config.LoadOptions) error {
lo.BaseEndpoint = endpoint
lo.Credentials = credentials.NewStaticCredentialsProvider(cfg.GetAWSKeyID(), cfg.GetAWSSecretKey(), cfg.GetAWSSessionToken())
lo.Region = cfg.GetAWSRegion()
return nil
})
require.NoError(t, err)
cfg.SetStoreClientFromS3Cfg(s3.NewFromConfig(acfg, func(o *s3.Options) {
o.UsePathStyle = true
}))
}
type PutObjectParams struct {
Key objectstore.BucketKey
File io.Reader
}
func PutObject(t testing.TB, ctx context.Context, cfg objectstore.ConfigProvider, params PutObjectParams) {
key := params.Key.String()
bucketName := GetBucketName(t)
_, err := cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
Bucket: &bucketName,
Key: &key,
Body: params.File,
})
require.NoError(t, err)
slog.Info("put object", "bucket", bucketName, "key", key)
}