From e7ab70be4990b5830f545f13301557208382f047 Mon Sep 17 00:00:00 2001
From: Brenna N Epp <brennae@google.com>
Date: Tue, 1 Nov 2022 23:01:29 +0000
Subject: [PATCH] test(storage): add gRPC error code to retry to fix flaky PAP
 test (#6974)

---
 storage/integration_test.go | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/storage/integration_test.go b/storage/integration_test.go
index bd5429c5d1d8..0d8e05d264cc 100644
--- a/storage/integration_test.go
+++ b/storage/integration_test.go
@@ -815,13 +815,16 @@ func TestIntegration_PublicAccessPrevention(t *testing.T) {
 
 		// Now, making object public or making bucket public should succeed. Run with
 		// retry because ACL settings may take time to propagate.
-		idempotentOrNilRetry := func(err error) bool {
-			return err == nil || ShouldRetry(err)
+		retrier := func(err error) bool {
+			// Once ACL settings propagate, PAP should no longer be enforced and the call will succeed.
+			// In the meantime, while PAP is enforced, trying to set ACL results in:
+			// 	-	FailedPrecondition for gRPC
+			// 	-	condition not met (412) for HTTP
+			return ShouldRetry(err) || status.Code(err) == codes.FailedPrecondition || extractErrCode(err) == http.StatusPreconditionFailed
 		}
 
 		ctxWithTimeout, cancelCtx := context.WithTimeout(ctx, time.Second*10)
-
-		a = o.Retryer(WithErrorFunc(idempotentOrNilRetry)).ACL()
+		a = o.Retryer(WithErrorFunc(retrier), WithPolicy(RetryAlways)).ACL()
 		err = a.Set(ctxWithTimeout, AllUsers, RoleReader)
 		cancelCtx()
 		if err != nil {
@@ -4872,3 +4875,16 @@ func skipGRPC(reason string) context.Context {
 func skipHTTP(reason string) context.Context {
 	return context.WithValue(context.Background(), skipTransportTestKey("http"), reason)
 }
+
+// Extract the error code if it's a googleapi.Error
+func extractErrCode(err error) int {
+	if err == nil {
+		return 0
+	}
+	var e *googleapi.Error
+	if errors.As(err, &e) {
+		return e.Code
+	}
+
+	return -1
+}