fleet/server/datastore/s3/testing_utils.go
Ian Littman 0d29f2bfc0
Add custom software icons (#32652)
For #29478, sans GitOps.

---------

Co-authored-by: RachelElysia <71795832+RachelElysia@users.noreply.github.com>
Co-authored-by: Konstantin Sykulev <konst@sykulev.com>
2025-09-05 17:31:03 -05:00

113 lines
3.1 KiB
Go

package s3
import (
"context"
"errors"
"os"
"testing"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/aws/aws-sdk-go-v2/service/s3/types"
"github.com/fleetdm/fleet/v4/server/config"
"github.com/stretchr/testify/require"
)
const (
accessKeyID = "minio"
secretAccessKey = "minio123!"
testEndpoint = "http://localhost:9000"
)
func SetupTestSoftwareInstallerStore(tb testing.TB, bucket, prefix string) *SoftwareInstallerStore {
store := setupTestStore(tb, bucket, prefix, NewSoftwareInstallerStore)
tb.Cleanup(func() { cleanupStore(tb, store.s3store) })
return store
}
func SetupTestBootstrapPackageStore(tb testing.TB, bucket, prefix string) *BootstrapPackageStore {
store := setupTestStore(tb, bucket, prefix, NewBootstrapPackageStore)
tb.Cleanup(func() { cleanupStore(tb, store.s3store) })
return store
}
func SetupTestSoftwareTitleIconStore(tb testing.TB, bucket, prefix string) *SoftwareTitleIconStore {
store := setupTestStore(tb, bucket, prefix, NewSoftwareTitleIconStore)
tb.Cleanup(func() { cleanupStore(tb, store.s3store) })
return store
}
type testBucketCreator interface {
CreateTestBucket(ctx context.Context, name string) error
}
func setupTestStore[T testBucketCreator](tb testing.TB, bucket, prefix string, newFn func(config.S3Config) (T, error)) T {
checkEnv(tb)
store, err := newFn(config.S3Config{
SoftwareInstallersBucket: bucket,
SoftwareInstallersPrefix: prefix,
SoftwareInstallersRegion: "minio",
SoftwareInstallersEndpointURL: testEndpoint,
SoftwareInstallersAccessKeyID: accessKeyID,
SoftwareInstallersSecretAccessKey: secretAccessKey,
SoftwareInstallersForceS3PathStyle: true,
SoftwareInstallersDisableSSL: true,
CarvesBucket: bucket,
CarvesPrefix: prefix,
CarvesRegion: "minio",
CarvesEndpointURL: testEndpoint,
CarvesAccessKeyID: accessKeyID,
CarvesSecretAccessKey: secretAccessKey,
CarvesForceS3PathStyle: true,
CarvesDisableSSL: true,
})
require.Nil(tb, err)
err = store.CreateTestBucket(context.Background(), bucket)
require.NoError(tb, err)
return store
}
func cleanupStore(tb testing.TB, store *s3store) {
checkEnv(tb)
ctx := context.Background()
resp, err := store.s3Client.ListObjects(ctx, &s3.ListObjectsInput{
Bucket: &store.bucket,
})
var noSuchBucket *types.NoSuchBucket
if errors.As(err, &noSuchBucket) {
// OK, nothing to clean-up if the bucket no longer exists, no error
return
}
require.NoError(tb, err)
var objs []types.ObjectIdentifier
for _, o := range resp.Contents {
objs = append(objs, types.ObjectIdentifier{
Key: o.Key,
})
}
if len(objs) > 0 {
_, err = store.s3Client.DeleteObjects(ctx, &s3.DeleteObjectsInput{
Bucket: &store.bucket,
Delete: &types.Delete{
Objects: objs,
},
})
require.NoError(tb, err)
}
_, err = store.s3Client.DeleteBucket(ctx, &s3.DeleteBucketInput{
Bucket: &store.bucket,
})
require.NoError(tb, err)
}
func checkEnv(tb testing.TB) {
if _, ok := os.LookupEnv("MINIO_STORAGE_TEST"); !ok {
tb.Skip("set MINIO_STORAGE_TEST environment variable to run S3-based tests")
}
}