Michael McGuinness
efe87321b2
Testing Tweaks * parallel * slowestquery * split * cleanup * health * dynamiccores * go * profile * timeout * commitmychanges * taskfile * sqlcheckstart * client * cost * ctxtimeout
91 lines
2.5 KiB
Go
91 lines
2.5 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, ctx context.Context, cfg objectstore.ConfigProvider) {
|
|
t.Helper()
|
|
bucketName := GetBucketName(t)
|
|
_, err := cfg.GetStoreClient().CreateBucket(ctx, &s3.CreateBucketInput{
|
|
Bucket: aws.String(bucketName),
|
|
})
|
|
require.NoError(t, err)
|
|
slog.Info("create bucket", "bucket", bucketName)
|
|
|
|
err = cfg.PingStoreByName(ctx, bucketName)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
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)
|
|
}
|