test(storage): RequesterPays test fix (#5722)

* test(storage): add retry on rate limit error for RequestorPays test
googleapis · Mar 10, 2022 · c2df3c1 · c2df3c1
1 parent 7f69bbb
commit c2df3c1
Showing 1 changed file with 35 additions and 33 deletions.
diff --git a/storage/integration_test.go b/storage/integration_test.go
@@ -2781,7 +2781,6 @@ func TestIntegration_RequesterPays(t *testing.T) {
 	// - (1a) must NOT have that permission on (3b).
 
 	ctx := context.Background()
-	h := testHelper{t}
 
 	// Start client with mainUserEmail creds
 	mainUserClient := testConfig(ctx, t)
@@ -2835,24 +2834,9 @@ func TestIntegration_RequesterPays(t *testing.T) {
 		return -1
 	}
 
-	bucketName := uidSpace.New()
-	requesterPaysBucket := mainUserClient.Bucket(bucketName)
-	bucketName2 := uidSpace.New()
-	requesterPaysBucket2 := mainUserClient.Bucket(bucketName2)
-
-	// Create a requester-pays bucket. The bucket is contained in the project mainProjectID
-	h.mustCreate(requesterPaysBucket, mainProjectID, &BucketAttrs{RequesterPays: true})
-	if err := requesterPaysBucket.ACL().Set(ctx, ACLEntity("user-"+otherUserEmail), RoleOwner); err != nil {
-		t.Fatalf("set ACL: %v", err)
-	}
-	defer h.mustDeleteBucket(requesterPaysBucket)
-
-	// Create a second requester-pays bucket to avoid rate limits.
-	h.mustCreate(requesterPaysBucket2, mainProjectID, &BucketAttrs{RequesterPays: true})
-	if err := requesterPaysBucket2.ACL().Set(ctx, ACLEntity("user-"+otherUserEmail), RoleOwner); err != nil {
-		t.Fatalf("set ACL: %v", err)
-	}
-	defer h.mustDeleteBucket(requesterPaysBucket2)
+	// We hit bucket rate limits with these test, so we retry
+	mainUserClient.SetRetry(WithPolicy(RetryAlways))
+	otherUserClient.SetRetry(WithPolicy(RetryAlways))
 
 	for _, test := range []struct {
 		desc          string
@@ -2895,6 +2879,9 @@ func TestIntegration_RequesterPays(t *testing.T) {
 	} {
 		t.Run(test.desc, func(t *testing.T) {
 			h := testHelper{t}
+			ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+			defer cancel()
+
 			printTestCase := func() string {
 				user := mainUserEmail
 				if test.client == otherUserClient {
@@ -2918,6 +2905,16 @@ func TestIntegration_RequesterPays(t *testing.T) {
 				}
 			}
 
+			bucketName := uidSpace.New()
+			requesterPaysBucket := mainUserClient.Bucket(bucketName)
+
+			// Create a requester-pays bucket. The bucket is contained in the project mainProjectID
+			h.mustCreate(requesterPaysBucket, mainProjectID, &BucketAttrs{RequesterPays: true})
+			if err := requesterPaysBucket.ACL().Set(ctx, ACLEntity("user-"+otherUserEmail), RoleOwner); err != nil {
+				t.Fatalf("set ACL: %v", err)
+			}
+			defer h.mustDeleteBucket(requesterPaysBucket)
+
 			// Make sure the object exists, so we don't get confused by ErrObjectNotExist.
 			// The later write we perform may fail so we always write to the object as the user
 			// with permissions on the containing bucket (mainUser).
@@ -2932,11 +2929,6 @@ func TestIntegration_RequesterPays(t *testing.T) {
 			if test.userProject != nil {
 				bucket = bucket.UserProject(*test.userProject)
 			}
-			// Create another bucket to avoid the per bucket update rate limit
-			bucket2 := test.client.Bucket(bucketName2)
-			if test.userProject != nil {
-				bucket2 = bucket2.UserProject(*test.userProject)
-			}
 
 			// Get bucket attrs
 			attrs, err := bucket.Attrs(ctx)
@@ -2960,10 +2952,10 @@ func TestIntegration_RequesterPays(t *testing.T) {
 			// We interleave buckets to get around the rate limit
 			entity := ACLEntity("domain-google.com")
 
-			checkforErrors("bucket acl set", bucket2.ACL().Set(ctx, entity, RoleReader))
+			checkforErrors("bucket acl set", bucket.ACL().Set(ctx, entity, RoleReader))
 			_, err = bucket.ACL().List(ctx)
 			checkforErrors("bucket acl list", err)
-			checkforErrors("bucket acl delete", bucket2.ACL().Delete(ctx, entity))
+			checkforErrors("bucket acl delete", bucket.ACL().Delete(ctx, entity))
 
 			// Object ACL operations
 			checkforErrors("object acl set", bucket.Object(objectName).ACL().Set(ctx, entity, RoleReader))
@@ -2974,15 +2966,31 @@ func TestIntegration_RequesterPays(t *testing.T) {
 			// Default object ACL operations
 			// Once again, we interleave buckets to avoid rate limits
 			checkforErrors("default object acl set", bucket.DefaultObjectACL().Set(ctx, entity, RoleReader))
-			_, err = bucket2.DefaultObjectACL().List(ctx)
+			_, err = bucket.DefaultObjectACL().List(ctx)
 			checkforErrors("default object acl list", err)
 			checkforErrors("default object acl delete", bucket.DefaultObjectACL().Delete(ctx, entity))
 
 			// Copy and compose
 			_, err = bucket.Object("copy").CopierFrom(bucket.Object(objectName)).Run(ctx)
 			checkforErrors("copy", err)
+			if err == nil {
+				// Delete created "copy" object if created successfully
+				defer func() {
+					if err := bucket.Object("copy").Delete(ctx); err != nil {
+						t.Fatalf("could not delete copy: %v", err)
+					}
+				}()
+			}
 			_, err = bucket.Object("compose").ComposerFrom(bucket.Object(objectName), bucket.Object("copy")).Run(ctx)
 			checkforErrors("compose", err)
+			if err == nil {
+				// Delete created "compose" object if created successfully
+				defer func() {
+					if err := bucket.Object("compose").Delete(ctx); err != nil {
+						t.Fatalf("could not delete compose: %v", err)
+					}
+				}()
+			}
 
 			// Delete object
 			if err = bucket.Object(objectName).Delete(ctx); err != nil {
@@ -2994,12 +3002,6 @@ func TestIntegration_RequesterPays(t *testing.T) {
 
 	}
 
-	// Clean up created objects for copy and compose
-	for _, obj := range []string{"copy", "compose"} {
-		if err := requesterPaysBucket.UserProject(mainProjectID).Object(obj).Delete(ctx); err != nil {
-			t.Fatalf("could not delete %q: %v", obj, err)
-		}
-	}
 }
 
 func TestIntegration_Notifications(t *testing.T) {