Improve ephemeral S3 buckets implementation for tests

Signed-off-by: Marko Mudrinić <[email protected]>

Author: xmudrii

tests/e2e/kubetest2-kops/aws/s3.go

@@ -32,46 +32,44 @@ import (
 	"k8s.io/klog/v2"
 )
 
-// We need to pick some region to query the AWS APIs through, even if we are not running on AWS.
+// defaultRegion is the region to query the AWS APIs through, this can be any AWS region is required even if we are not
+// running on AWS.
 const defaultRegion = "us-east-2"
 
-// It contains S3Client, an Amazon S3 service client that is used to perform bucket
-// and object actions.
-type awsClient struct {
-	S3Client *s3.Client
+// Client contains S3 and STS clients that are used to perform bucket and object actions.
+type Client struct {
+	s3Client  *s3.Client
+	stsClient *sts.Client
 }
 
-func NewAWSClient(ctx context.Context) (*awsClient, error) {
+// NewAWSClient returns a new instance of awsClient configured to work in the default region (us-east-2).
+func NewClient(ctx context.Context) (*Client, error) {
 	cfg, err := awsconfig.LoadDefaultConfig(ctx,
 		awsconfig.WithRegion(defaultRegion))
 	if err != nil {
-		return nil, fmt.Errorf("failed to load AWS config: %w", err)
+		return nil, fmt.Errorf("loading AWS config: %w", err)
 	}
-	return &awsClient{
-		S3Client: s3.NewFromConfig(cfg),
+
+	return &Client{
+		s3Client:  s3.NewFromConfig(cfg),
+		stsClient: sts.NewFromConfig(cfg),
 	}, nil
 }
 
-// AWSBucketName constructs an unique bucket name using the AWS account ID on region us-east-2
-func AWSBucketName(ctx context.Context) (string, error) {
-	config, err := awsconfig.LoadDefaultConfig(ctx, awsconfig.WithRegion(string(types.BucketLocationConstraintUsEast2)))
-	if err != nil {
-		return "", fmt.Errorf("failed to load AWS config: %w", err)
-	}
-
-	stsSvc := sts.NewFromConfig(config)
-	callerIdentity, err := stsSvc.GetCallerIdentity(ctx, &sts.GetCallerIdentityInput{})
+// BucketName constructs an unique bucket name using the AWS account ID in the default region (us-east-2).
+func (c Client) BucketName(ctx context.Context) (string, error) {
+	callerIdentity, err := c.stsClient.GetCallerIdentity(ctx, &sts.GetCallerIdentityInput{})
 	if err != nil {
 		return "", fmt.Errorf("building AWS STS presigned request: %w", err)
 	}
 
-	// Add timestamp suffix
-	timestamp := time.Now().Format("01022006")
-	bucket := fmt.Sprintf("k8s-infra-kops-%s", *callerIdentity.Account)
-	bucket = fmt.Sprintf("%s-%s", bucket, timestamp)
+	// Construct the bucket name based on the AWS account ID and the current timestamp
+	timestamp := time.Now().Format("20060102150405")
+	bucket := fmt.Sprintf("k8s-infra-kops-%s-%s", *callerIdentity.Account, timestamp)
 
 	bucket = strings.ToLower(bucket)
-	bucket = regexp.MustCompile("[^a-z0-9-]").ReplaceAllString(bucket, "") // Only allow lowercase letters, numbers, and hyphens
+	// Only allow lowercase letters, numbers, and hyphens
+	bucket = regexp.MustCompile("[^a-z0-9-]").ReplaceAllString(bucket, "")
 
 	if len(bucket) > 63 {
 		bucket = bucket[:63] // Max length is 63
@@ -80,73 +78,88 @@ func AWSBucketName(ctx context.Context) (string, error) {
 	return bucket, nil
 }
 
-func (client awsClient) EnsureS3Bucket(ctx context.Context, bucketName string, publicRead bool) error {
-	_, err := client.S3Client.CreateBucket(ctx, &s3.CreateBucketInput{
+// EnsureS3Bucket creates a new S3 bucket with the given name and public read permissions.
+func (c Client) EnsureS3Bucket(ctx context.Context, bucketName string, publicRead bool) error {
+	_, err := c.s3Client.CreateBucket(ctx, &s3.CreateBucketInput{
 		Bucket: aws.String(bucketName),
 		CreateBucketConfiguration: &types.CreateBucketConfiguration{
-			LocationConstraint: types.BucketLocationConstraintUsEast2,
+			LocationConstraint: defaultRegion,
 		},
 	})
-
-	var exists *types.BucketAlreadyExists
 	if err != nil {
+		var exists *types.BucketAlreadyExists
 		if errors.As(err, &exists) {
-			klog.Infof("Bucket %s already exists.\n", bucketName)
-			err = exists
-		}
-	} else {
-		err := s3.NewBucketExistsWaiter(client.S3Client).Wait(
-			ctx, &s3.HeadBucketInput{
-				Bucket: aws.String(bucketName),
-			},
-			time.Minute)
-		if err != nil {
-			klog.Infof("Failed attempt to wait for bucket %s to exist.", bucketName)
+			klog.Infof("Bucket %s already exists\n", bucketName)
+		} else {
+			klog.Infof("Error creating bucket %s, err: %v\n", bucketName, err)
 		}
+
+		return fmt.Errorf("creating bucket %s: %w", bucketName, err)
+	}
+
+	// Wait for the bucket to be created
+	err = s3.NewBucketExistsWaiter(c.s3Client).Wait(
+		ctx, &s3.HeadBucketInput{
+			Bucket: aws.String(bucketName),
+		},
+		time.Minute)
+	if err != nil {
+		klog.Infof("Failed attempt to wait for bucket %s to exist, err: %v", bucketName, err)
+
+		return fmt.Errorf("waiting for bucket %s to exist: %w", bucketName, err)
 	}
 
 	klog.Infof("Bucket %s created successfully", bucketName)
 
-	if err == nil && publicRead {
-		err = client.setPublicReadPolicy(ctx, bucketName)
+	if publicRead {
+		err = c.setPublicReadPolicy(ctx, bucketName)
 		if err != nil {
-			klog.Errorf("Failed to set public read policy on bucket %s: %v", bucketName, err)
-			return err
+			klog.Errorf("Failed to set public read policy on bucket %s, err: %v", bucketName, err)
+
+			return fmt.Errorf("setting public read policy for bucket %s: %w", bucketName, err)
 		}
+
 		klog.Infof("Public read policy set on bucket %s", bucketName)
 	}
 
-	return err
+	return nil
 }
 
-func (client awsClient) DeleteS3Bucket(ctx context.Context, bucketName string) error {
-	_, err := client.S3Client.DeleteBucket(ctx, &s3.DeleteBucketInput{
+// DeleteS3Bucket deletes a S3 bucket with the given name.
+func (c Client) DeleteS3Bucket(ctx context.Context, bucketName string) error {
+	_, err := c.s3Client.DeleteBucket(ctx, &s3.DeleteBucketInput{
 		Bucket: aws.String(bucketName),
 	})
 	if err != nil {
 		var noBucket *types.NoSuchBucket
 		if errors.As(err, &noBucket) {
-			klog.Infof("Bucket %s does not exits", bucketName)
-			err = noBucket
-		} else {
-			klog.Infof("Couldn't delete bucket %s. Reason: %v", bucketName, err)
-		}
-	} else {
-		err = s3.NewBucketNotExistsWaiter(client.S3Client).Wait(
-			ctx, &s3.HeadBucketInput{
-				Bucket: aws.String(bucketName),
-			},
-			time.Minute)
-		if err != nil {
-			klog.Infof("Failed attempt to wait for bucket %s to be deleted", bucketName)
+			klog.Infof("Bucket %s does not exits.", bucketName)
+
+			return nil
 		} else {
-			klog.Infof("Bucket %s deleted", bucketName)
+			klog.Infof("Couldn't delete bucket %s, err: %v", bucketName, err)
+
+			return fmt.Errorf("deleting bucket %s: %w", bucketName, err)
 		}
 	}
-	return err
+
+	err = s3.NewBucketNotExistsWaiter(c.s3Client).Wait(
+		ctx, &s3.HeadBucketInput{
+			Bucket: aws.String(bucketName),
+		},
+		time.Minute)
+	if err != nil {
+		klog.Infof("Failed attempt to wait for bucket %s to be deleted, err: %v", bucketName, err)
+
+		return fmt.Errorf("waiting for bucket %s to be deleted, err: %w", bucketName, err)
+	}
+
+	klog.Infof("Bucket %s deleted", bucketName)
+
+	return nil
 }
 
-func (client awsClient) setPublicReadPolicy(ctx context.Context, bucketName string) error {
+func (c Client) setPublicReadPolicy(ctx context.Context, bucketName string) error {
 	policy := fmt.Sprintf(`{
     "Version": "2012-10-17",
     "Statement": [
@@ -160,13 +173,13 @@ func (client awsClient) setPublicReadPolicy(ctx context.Context, bucketName stri
     ]
   }`, bucketName)
 
-	_, err := client.S3Client.PutBucketPolicy(ctx, &s3.PutBucketPolicyInput{
+	_, err := c.s3Client.PutBucketPolicy(ctx, &s3.PutBucketPolicyInput{
 		Bucket: aws.String(bucketName),
 		Policy: aws.String(policy),
 	})
 	if err != nil {
-		return fmt.Errorf("failed to put bucket policy for %s: %w", bucketName, err)
+		return fmt.Errorf("putting bucket policy for %s: %w", bucketName, err)
 	}
 
-	return err
+	return nil
 }

tests/e2e/kubetest2-kops/deployer/common.go

@@ -53,6 +53,12 @@ func (d *deployer) initialize() error {
 
 	switch d.CloudProvider {
 	case "aws":
+		client, err := aws.NewClient(context.Background())
+		if err != nil {
+			return fmt.Errorf("init failed to build AWS client: %w", err)
+		}
+		d.aws = client
+
 		if d.SSHPrivateKeyPath == "" {
 			d.SSHPrivateKeyPath = os.Getenv("AWS_SSH_PRIVATE_KEY_FILE")
 		}
@@ -318,23 +324,20 @@ func defaultClusterName(cloudProvider string) (string, error) {
 // stateStore returns the kops state store to use
 // defaulting to values used in prow jobs
 func (d *deployer) stateStore() string {
+	if d.stateStoreName != "" {
+		return d.stateStoreName
+	}
 	ss := os.Getenv("KOPS_STATE_STORE")
 	if ss == "" {
 		switch d.CloudProvider {
 		case "aws":
-			ctx := context.TODO()
-			bucketName, err := aws.AWSBucketName(ctx)
+			ctx := context.Background()
+			bucketName, err := d.aws.BucketName(ctx)
 			if err != nil {
 				klog.Fatalf("Failed to generate bucket name: %v", err)
-			}
-			awsClient, err := aws.NewAWSClient(ctx)
-			if err != nil {
-				klog.Fatalf("failed to load client config: %v", err)
-			}
-			if err := awsClient.EnsureS3Bucket(ctx, bucketName, d.PublicReadOnlyBucket); err != nil {
-				klog.Fatalf("Failed to ensure S3 bucket exists: %v", err)
 				return ""
 			}
+			d.createBucket = true
 			ss = "s3://" + bucketName
 		case "gce":
 			d.createBucket = true
@@ -343,22 +346,31 @@ func (d *deployer) stateStore() string {
 			ss = "do://e2e-kops-space"
 		}
 	}
+
+	d.stateStoreName = ss
 	return ss
 }
 
 // discoveryStore returns the VFS path to use for public OIDC documents
 func (d *deployer) discoveryStore() string {
+	if d.discoveryStoreName != "" {
+		return d.discoveryStoreName
+	}
 	discovery := os.Getenv("KOPS_DISCOVERY_STORE")
 	if discovery == "" {
 		switch d.CloudProvider {
 		case "aws":
 			discovery = "s3://k8s-kops-ci-prow"
 		}
 	}
+	d.discoveryStoreName = discovery
 	return discovery
 }
 
 func (d *deployer) stagingStore() string {
+	if d.stagingStoreName != "" {
+		return d.stagingStoreName
+	}
 	sb := os.Getenv("KOPS_STAGING_BUCKET")
 	if sb == "" {
 		switch d.CloudProvider {
@@ -367,6 +379,7 @@ func (d *deployer) stagingStore() string {
 			sb = "gs://" + gce.GCSBucketName(d.GCPProject, "staging")
 		}
 	}
+	d.stagingStoreName = sb
 	return sb
 }
 

tests/e2e/kubetest2-kops/deployer/deployer.go

@@ -25,6 +25,7 @@ import (
 	"github.com/octago/sflags/gen/gpflag"
 	"github.com/spf13/pflag"
 	"k8s.io/klog/v2"
+	"k8s.io/kops/tests/e2e/kubetest2-kops/aws"
 	"k8s.io/kops/tests/e2e/kubetest2-kops/builder"
 	"k8s.io/kops/tests/e2e/pkg/target"
 
@@ -57,8 +58,6 @@ type deployer struct {
 	CreateArgs             string   `flag:"create-args" desc:"Extra space-separated arguments passed to 'kops create cluster'"`
 	KopsBinaryPath         string   `flag:"kops-binary-path" desc:"The path to kops executable used for testing"`
 	KubernetesFeatureGates string   `flag:"kubernetes-feature-gates" desc:"Feature Gates to enable on Kubernetes components"`
-	createBucket           bool     `flag:"-"`
-	PublicReadOnlyBucket   bool     `flag:"-"`
 
 	// ControlPlaneCount specifies the number of VMs in the control-plane.
 	ControlPlaneCount int `flag:"control-plane-count" desc:"Number of control-plane instances"`
@@ -91,6 +90,13 @@ type deployer struct {
 	manifestPath string
 	terraform    *target.Terraform
 
+	aws *aws.Client
+
+	createBucket       bool
+	stateStoreName     string
+	discoveryStoreName string
+	stagingStoreName   string
+
 	// boskos struct field will be non-nil when the deployer is
 	// using boskos to acquire a GCP project
 	boskos                  *client.Client

tests/e2e/kubetest2-kops/deployer/down.go

@@ -17,6 +17,7 @@ limitations under the License.
 package deployer
 
 import (
+	"context"
 	"fmt"
 	"strings"
 
@@ -72,9 +73,17 @@ func (d *deployer) Down() error {
 		return err
 	}
 
-	if d.CloudProvider == "gce" && d.createBucket {
-		gce.DeleteGCSBucket(d.stateStore(), d.GCPProject)
-		gce.DeleteGCSBucket(d.stagingStore(), d.GCPProject)
+	if d.createBucket {
+		switch d.CloudProvider {
+		case "aws":
+			ctx := context.Background()
+			if err := d.aws.DeleteS3Bucket(ctx, d.stateStore()); err != nil {
+				return err
+			}
+		case "gce":
+			gce.DeleteGCSBucket(d.stateStore(), d.GCPProject)
+			gce.DeleteGCSBucket(d.stagingStore(), d.GCPProject)
+		}
 	}
 
 	if d.boskos != nil {

tests/e2e/kubetest2-kops/deployer/up.go

@@ -17,6 +17,7 @@ limitations under the License.
 package deployer
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	osexec "os/exec"
@@ -61,9 +62,17 @@ func (d *deployer) Up() error {
 		_ = d.Down()
 	}
 
-	if d.CloudProvider == "gce" && d.createBucket {
-		if err := gce.EnsureGCSBucket(d.stateStore(), d.GCPProject, false); err != nil {
-			return err
+	if d.createBucket {
+		switch d.CloudProvider {
+		case "aws":
+			ctx := context.Background()
+			if err := d.aws.EnsureS3Bucket(ctx, d.stateStore(), false); err != nil {
+				return err
+			}
+		case "gce":
+			if err := gce.EnsureGCSBucket(d.stateStore(), d.GCPProject, false); err != nil {
+				return err
+			}
 		}
 	}