Support terminal status on initial op in sync poller (#23564)

* Support terminal status on initial op in sync poller * Add test for default sync poller fix * Unused imports
Azure · Aug 17, 2021 · 72c87fe · 72c87fe
1 parent 7bc4506
commit 72c87fe
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 0 deletions.
diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/util/polling/DefaultSyncPoller.java b/sdk/core/azure-core/src/main/java/com/azure/core/util/polling/DefaultSyncPoller.java
@@ -72,6 +72,9 @@ final class DefaultSyncPoller<T, U> implements SyncPoller<T, U> {
         //
         this.pollingContext.setOnetimeActivationResponse(this.activationResponse);
         this.pollingContext.setLatestResponse(this.activationResponse);
+        if (this.activationResponse.getStatus().isComplete()) {
+            this.terminalPollContext = this.pollingContext;
+        }
     }
 
     @Override

diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/util/polling/PollerTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/util/polling/PollerTests.java
@@ -32,6 +32,7 @@
 import static com.azure.core.util.polling.PollerFlux.error;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.when;
 
@@ -213,6 +214,36 @@ public void noPollingForSynchronouslyCompletedActivationTest() {
         assertEquals(1, activationCallCount[0]);
     }
 
+    @Test
+    public void noPollingForSynchronouslyCompletedActivationInSyncPollerTest() {
+        int[] activationCallCount = new int[1];
+        activationCallCount[0] = 0;
+        when(activationOperationWithResponse.apply(any())).thenReturn(Mono.defer(() -> {
+            activationCallCount[0]++;
+            return Mono.just(new PollResponse<Response>(LongRunningOperationStatus.SUCCESSFULLY_COMPLETED,
+                new Response("ActivationDone")));
+        }));
+
+        SyncPoller<Response, CertificateOutput> syncPoller = create(
+            Duration.ofMillis(10),
+            activationOperationWithResponse,
+            pollOperation,
+            cancelOperation,
+            fetchResultOperation)
+            .getSyncPoller();
+
+        when(pollOperation.apply(any())).thenReturn(
+            Mono.error(new RuntimeException("Polling shouldn't happen for synchronously completed activation.")));
+
+        try {
+            PollResponse<Response> response = syncPoller.waitForCompletion(Duration.ofSeconds(1));
+            assertEquals(LongRunningOperationStatus.SUCCESSFULLY_COMPLETED, response.getStatus());
+            assertEquals(1, activationCallCount[0]);
+        } catch (Exception e) {
+            fail("SyncPoller did not complete on activation", e);
+        }
+    }
+
     @Test
     public void ensurePollingForInProgressActivationResponseTest() {
         final Duration retryAfter = Duration.ofMillis(10);