core, NPE protect on Context in PollingStrategy (Azure#28507)

* core, NPE protect on Context in PollingStrategy

* add test case for Context passed to PollingStragety

Author: weidongxu-microsoft

sdk/core/azure-core/src/main/java/com/azure/core/util/polling/LocationPollingStrategy.java

@@ -75,7 +75,7 @@ public LocationPollingStrategy(HttpPipeline httpPipeline, ObjectSerializer seria
     public LocationPollingStrategy(HttpPipeline httpPipeline, ObjectSerializer serializer, Context context) {
         this.httpPipeline = Objects.requireNonNull(httpPipeline, "'httpPipeline' cannot be null");
         this.serializer = (serializer == null) ? DEFAULT_SERIALIZER : serializer;
-        this.context = context;
+        this.context = context == null ? Context.NONE : context;
     }
 
     @Override

sdk/core/azure-core/src/main/java/com/azure/core/util/polling/OperationResourcePollingStrategy.java

@@ -74,7 +74,7 @@ public OperationResourcePollingStrategy(HttpPipeline httpPipeline, ObjectSeriali
         this.serializer = serializer != null ? serializer : new DefaultJsonSerializer();
         this.operationLocationHeaderName = operationLocationHeaderName != null ? operationLocationHeaderName
             : DEFAULT_OPERATION_LOCATION_HEADER;
-        this.context = context;
+        this.context = context == null ? Context.NONE : context;
     }
 
 

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

@@ -11,10 +11,13 @@
 import com.azure.core.http.MockHttpResponse;
 import com.azure.core.http.rest.Response;
 import com.azure.core.http.rest.SimpleResponse;
+import com.azure.core.util.Context;
 import com.azure.core.util.serializer.TypeReference;
 import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
@@ -302,6 +305,68 @@ public void locationPollingStrategySucceedsOnPollWithPostLocationHeader() {
         assertEquals(1, activationCallCount[0]);
     }
 
+    @Test
+    public void pollingStrategyPassContextToHttpClient() {
+        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("Location", mockPollUrl),
+                new PollResult("InProgress"));
+            return Mono.just(response);
+        }));
+        HttpRequest pollRequest = new HttpRequest(HttpMethod.GET, mockPollUrl);
+        ArgumentCaptor<Context> contextArgument = ArgumentCaptor.forClass(Context.class);
+        when(httpClient.send(any(), contextArgument.capture()))
+            .thenAnswer(iom -> {
+                HttpRequest req = iom.getArgument(0);
+                if (mockPollUrl.equals(req.getUrl().toString())) {
+                    return Mono.just(new MockHttpResponse(pollRequest, 200,
+                        new HttpHeaders().set("Location", finalResultUrl),
+                        new PollResult("Succeeded")));
+                } else if (finalResultUrl.equals(req.getUrl().toString())) {
+                    return Mono.just(new MockHttpResponse(pollRequest, 200, new HttpHeaders(),
+                        new PollResult("final-state")));
+                } else {
+                    return Mono.error(new IllegalArgumentException("Unknown request URL " + req.getUrl()));
+                }
+            });
+
+        // PollingStrategy with context = Context.NONE
+        PollerFlux<PollResult, PollResult> pollerFlux = PollerFlux.create(
+            Duration.ofSeconds(1),
+            () -> activationOperation.get(),
+            new DefaultPollingStrategy<>(new HttpPipelineBuilder().httpClient(httpClient).build(), null, null),
+            new TypeReference<PollResult>() { }, new TypeReference<PollResult>() { });
+
+        StepVerifier.create(pollerFlux.map(AsyncPollResponse::getStatus))
+            .expectSubscription()
+            .expectNext(LongRunningOperationStatus.SUCCESSFULLY_COMPLETED)
+            .verifyComplete();
+        Assertions.assertEquals(Context.NONE, contextArgument.getValue());
+
+        // PollingStrategy with context
+        final Context context = new Context("key", "value");
+        pollerFlux = PollerFlux.create(
+            Duration.ofSeconds(1),
+            () -> activationOperation.get(),
+            new DefaultPollingStrategy<>(new HttpPipelineBuilder().httpClient(httpClient).build(), null, context),
+            new TypeReference<PollResult>() { }, new TypeReference<PollResult>() { });
+
+        StepVerifier.create(pollerFlux.map(AsyncPollResponse::getStatus))
+            .expectSubscription()
+            .expectNext(LongRunningOperationStatus.SUCCESSFULLY_COMPLETED)
+            .verifyComplete();
+        Assertions.assertEquals(context, contextArgument.getValue());
+
+        assertEquals(2, activationCallCount[0]);
+    }
+
     public static class PollResult {
         private String status;
         private String resourceLocation;