fix flaky test (#22622)

Azure · Jun 28, 2021 · ff74916 · ff74916
1 parent f611436
commit ff74916
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/...orage-common/src/test-shared/java/com/azure/storage/common/test/shared/StorageSpec.groovy b/...orage-common/src/test-shared/java/com/azure/storage/common/test/shared/StorageSpec.groovy
@@ -13,6 +13,10 @@ import com.azure.core.util.ServiceVersion
 import com.azure.identity.EnvironmentCredentialBuilder
 import spock.lang.Specification
 
+import java.time.Duration
+import java.util.function.Predicate
+import java.util.function.Supplier
+
 class StorageSpec extends Specification {
     private static final TestEnvironment ENVIRONMENT = TestEnvironment.getInstance()
     private static final HttpClient HTTP_CLIENT = new NettyAsyncHttpClientBuilder().build()
@@ -87,4 +91,20 @@ class StorageSpec extends Specification {
             .map { it.getToken() }
             .block()
     }
+
+    protected <T, E extends Exception> T retry(
+        Supplier<T> action, Predicate<E> retryPredicate,
+        int times=6, Duration delay=Duration.ofSeconds(10)) {
+        for (i in 0..<times) {
+            try {
+                return action.get()
+            } catch (Exception e) {
+                if (!retryPredicate(e)) {
+                    throw e
+                } else {
+                    Thread.sleep(delay.toMillis())
+                }
+            }
+        }
+    }
 }
diff --git a/...age-file-share/src/test/java/com/azure/storage/file/share/specialized/LeaseAPITest.groovy b/...age-file-share/src/test/java/com/azure/storage/file/share/specialized/LeaseAPITest.groovy
@@ -3,13 +3,15 @@
 
 package com.azure.storage.file.share.specialized
 
+import com.azure.core.http.rest.Response
 import com.azure.storage.common.test.shared.extensions.RequiredServiceVersion
 import com.azure.storage.file.share.APISpec
 import com.azure.storage.file.share.ShareClient
 import com.azure.storage.file.share.ShareFileClient
 import com.azure.storage.file.share.ShareServiceVersion
 import com.azure.storage.file.share.models.LeaseDurationType
 import com.azure.storage.file.share.models.LeaseStateType
+import com.azure.storage.file.share.models.ShareErrorCode
 import com.azure.storage.file.share.models.ShareStorageException
 import com.azure.storage.file.share.options.ShareAcquireLeaseOptions
 import com.azure.storage.file.share.options.ShareBreakLeaseOptions
@@ -203,9 +205,13 @@ class LeaseAPITest extends APISpec {
         setup:
         def shareSnapshot = shareClient.createSnapshot().getSnapshot()
         def shareClient = shareBuilderHelper(shareClient.getShareName(), shareSnapshot).buildClient()
+        def leaseClient = createLeaseClient(shareClient)
 
         when:
-        def resp = createLeaseClient(shareClient).acquireLeaseWithResponse(new ShareAcquireLeaseOptions().setDuration(-1), null, null)
+        def resp = retry({
+            leaseClient
+                .acquireLeaseWithResponse(new ShareAcquireLeaseOptions().setDuration(-1), null, null)
+        },{ShareStorageException it -> it.errorCode == ShareErrorCode.SHARE_SNAPSHOT_IN_PROGRESS})
 
         then:
         resp.getStatusCode() == 201