Validate LRO Retry behavior for failure service responses (Azure#31725)

Author: samvaity

sdk/core/azure-core/src/test/java/com/azure/core/util/polling/PollingStrategyTests.java

@@ -6,9 +6,11 @@
 import com.azure.core.http.HttpClient;
 import com.azure.core.http.HttpHeaders;
 import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpPipeline;
 import com.azure.core.http.HttpPipelineBuilder;
 import com.azure.core.http.HttpRequest;
 import com.azure.core.http.MockHttpResponse;
+import com.azure.core.http.policy.RetryPolicy;
 import com.azure.core.http.rest.Response;
 import com.azure.core.http.rest.SimpleResponse;
 import com.azure.core.implementation.serializer.DefaultJsonSerializer;
@@ -19,6 +21,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
@@ -28,7 +31,9 @@
 import reactor.test.StepVerifier;
 
 import java.time.Duration;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
@@ -492,6 +497,64 @@ public void pollingStrategyPassContextToHttpClient() {
         assertEquals(3, activationCallCount[0]);
     }
 
+    @ParameterizedTest
+    @MethodSource("statusCodeProvider")
+    public void retryPollingOperationWithPostActivationOperation(int[] args) {
+        int[] activationCallCount = new int[1];
+        activationCallCount[0] = 0;
+        String mockPollUrl = "http://localhost/poll";
+        String finalResultUrl = "http://localhost/final";
+        when(activationOperation.get()).thenReturn(Mono.defer(() -> {
+            activationCallCount[0]++;
+            SimpleResponse<PollResult> response = new SimpleResponse<>(
+                new HttpRequest(HttpMethod.POST, "http://localhost"),
+                200,
+                new HttpHeaders().set("Operation-Location", mockPollUrl).set("Location", finalResultUrl),
+                new PollResult("InProgress"));
+            return Mono.just(response);
+        }));
+
+        HttpRequest pollRequest = new HttpRequest(HttpMethod.GET, mockPollUrl);
+        AtomicInteger attemptCount = new AtomicInteger();
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .policies(new RetryPolicy())
+            .httpClient(request -> {
+                int count = attemptCount.getAndIncrement();
+                if (mockPollUrl.equals(request.getUrl().toString()) && count == 0) {
+                    return Mono.just(new MockHttpResponse(pollRequest, args[0],
+                        new HttpHeaders().set("Location", finalResultUrl),
+                        new PollResult("Retry")));
+                } else if (mockPollUrl.equals(request.getUrl().toString()) && count == 1) {
+                    return Mono.just(new MockHttpResponse(pollRequest, args[1],
+                        new HttpHeaders().set("Location", finalResultUrl),
+                        new PollResult("Succeeded")));
+                } else if (finalResultUrl.equals(request.getUrl().toString())) {
+                    return Mono.just(new MockHttpResponse(pollRequest, args[2], new HttpHeaders(),
+                        new PollResult("final-state")));
+                } else {
+                    return Mono.error(new IllegalArgumentException("Unknown request URL " + request.getUrl()));
+                }
+            })
+            .build();
+        PollerFlux<PollResult, PollResult> pollerFlux = PollerFlux.create(
+            Duration.ofSeconds(1),
+            () -> activationOperation.get(),
+            new OperationResourcePollingStrategy<>(pipeline),
+            new TypeReference<PollResult>() { }, new TypeReference<PollResult>() { });
+
+        StepVerifier.create(pollerFlux.takeUntil(apr -> apr.getStatus().isComplete()).last().flatMap(AsyncPollResponse::getFinalResult))
+            .expectNextMatches(pollResult -> "final-state".equals(pollResult.getStatus()))
+            .verifyComplete();
+        assertEquals(args[3], attemptCount.get());
+        assertEquals(1, activationCallCount[0]);
+    }
+
+    static Stream<int[]> statusCodeProvider() {
+        return Stream.of(
+            new int[]{500, 200, 200, 3},
+            new int[]{200, 500, 200, 2});
+    }
+
     public static class PollResult {
         private String status;
         private String resourceLocation;