[core-lro] merge azureAsync and location strategies

sdk/core/core-lro/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Fix polling so that resources created in a different URL will be retrieved once polling is done. [PR #20656](https://github.com/Azure/azure-sdk-for-js/pull/20656)
+
 ### Other Changes
 
 ## 2.2.3 (2022-01-06)

sdk/core/core-lro/package.json

@@ -75,7 +75,7 @@
     "clean": "rimraf dist dist-* types *.log browser statistics.html coverage src/**/*.js test/**/*.js",
     "execute:samples": "echo skipped",
     "extract-api": "tsc -p . && api-extractor run --local",
-    "format": "prettier --write --config ../../../.prettierrc.json --ignore-path ../../../.prettierignore \"src/**/*.ts\" \"test/**/*.ts\" \"samples-dev/**/*.ts\" \"*.{js,json}\"",
+    "format": "prettier --write --config ../../../.prettierrc.json --ignore-path ../../../.prettierignore \"src/**/*.ts\" \"test/**/*.ts\" \"*.{js,json}\"",
     "integration-test:browser": "echo skipped",
     "integration-test:node": "echo skipped",
     "integration-test": "npm run integration-test:node && npm run integration-test:browser",
@@ -101,6 +101,7 @@
     "@azure/core-rest-pipeline": "^1.1.0",
     "@azure/eslint-plugin-azure-sdk": "^3.0.0",
     "@azure/dev-tool": "^1.0.0",
+    "@azure/test-utils": "^1.0.0",
     "@microsoft/api-extractor": "^7.18.11",
     "@types/chai": "^4.1.6",
     "@types/mocha": "^7.0.2",

sdk/core/core-lro/src/lroEngine/locationPolling.ts

@@ -1,18 +1,79 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { LroResponse, LroStatus, RawResponse } from "./models";
+import {
+  LongRunningOperation,
+  LroBody,
+  LroResourceLocationConfig,
+  LroResponse,
+  LroStatus,
+  RawResponse,
+  failureStates,
+  successStates,
+} from "./models";
 import { isUnexpectedPollingResponse } from "./requestUtils";
 
-function isLocationPollingDone(rawResponse: RawResponse): boolean {
-  return !isUnexpectedPollingResponse(rawResponse) && rawResponse.statusCode !== 202;
+function isPollingDone(rawResponse: RawResponse): boolean {
+  if (isUnexpectedPollingResponse(rawResponse) || rawResponse.statusCode === 202) {
+    return false;
+  }
+  const { status } = (rawResponse.body as LroBody) ?? {};
+  const state = typeof status === "string" ? status.toLowerCase() : "succeeded";
+  if (isUnexpectedPollingResponse(rawResponse) || failureStates.includes(state)) {
+    throw new Error(`The long running operation has failed. The provisioning state: ${state}.`);
+  }
+  return successStates.includes(state);
+}
+
+/**
+ * Sends a request to the URI of the provisioned resource if needed.
+ */
+async function sendFinalRequest<TResult>(
+  lro: LongRunningOperation<TResult>,
+  resourceLocation: string,
+  lroResourceLocationConfig?: LroResourceLocationConfig
+): Promise<LroResponse<TResult> | undefined> {
+  switch (lroResourceLocationConfig) {
+    case "original-uri":
+      return lro.sendPollRequest(lro.requestPath);
+    case "azure-async-operation":
+      return undefined;
+    case "location":
+    default:
+      return lro.sendPollRequest(resourceLocation ?? lro.requestPath);
+  }
 }
 
 export function processLocationPollingOperationResult<TResult>(
-  response: LroResponse<TResult>
-): LroStatus<TResult> {
-  return {
-    ...response,
-    done: isLocationPollingDone(response.rawResponse),
+  lro: LongRunningOperation<TResult>,
+  resourceLocation?: string,
+  lroResourceLocationConfig?: LroResourceLocationConfig
+): (response: LroResponse<TResult>) => LroStatus<TResult> {
+  return (response: LroResponse<TResult>): LroStatus<TResult> => {
+    if (isPollingDone(response.rawResponse)) {
+      if (resourceLocation === undefined) {
+        return { ...response, done: true };
+      } else {
+        return {
+          ...response,
+          done: false,
+          next: async () => {
+            const finalResponse = await sendFinalRequest(
+              lro,
+              resourceLocation,
+              lroResourceLocationConfig
+            );
+            return {
+              ...(finalResponse ?? response),
+              done: true,
+            };
+          },
+        };
+      }
+    }
+    return {
+      ...response,
+      done: false,
+    };
   };
 }

sdk/core/core-lro/src/lroEngine/models.ts

@@ -83,7 +83,7 @@ export interface LroResponse<T> {
 }
 
 /** The type of which LRO implementation being followed by a specific API. */
-export type LroMode = "AzureAsync" | "Location" | "Body";
+export type LroMode = "Location" | "Body";
 
 /**
  * The configuration of a LRO to determine how to perform polling and checking whether the operation has completed.

sdk/core/core-lro/src/lroEngine/requestUtils.ts

@@ -12,8 +12,8 @@ import { LroConfig, RawResponse } from "./models";
 export function getPollingUrl(rawResponse: RawResponse, defaultPath: string): string {
   return (
     getAzureAsyncOperation(rawResponse) ??
-    getLocation(rawResponse) ??
     getOperationLocation(rawResponse) ??
+    getLocation(rawResponse) ??
     defaultPath
   );
 }
@@ -30,25 +30,39 @@ function getAzureAsyncOperation(rawResponse: RawResponse): string | undefined {
   return rawResponse.headers["azure-asyncoperation"];
 }
 
+function findResourceLocation(
+  requestMethod: string,
+  rawResponse: RawResponse,
+  requestPath: string
+): string | undefined {
+  switch (requestMethod) {
+    case "PUT": {
+      return requestPath;
+    }
+    case "POST":
+    case "PATCH": {
+      return getLocation(rawResponse);
+    }
+    default: {
+      return undefined;
+    }
+  }
+}
+
 export function inferLroMode(
   requestPath: string,
   requestMethod: string,
   rawResponse: RawResponse
 ): LroConfig {
-  if (getAzureAsyncOperation(rawResponse) !== undefined) {
-    return {
-      mode: "AzureAsync",
-      resourceLocation:
-        requestMethod === "PUT"
-          ? requestPath
-          : requestMethod === "POST" || requestMethod === "PATCH"
-          ? getLocation(rawResponse)
-          : undefined,
-    };
-  } else if (
-    getLocation(rawResponse) !== undefined ||
+  if (
+    getAzureAsyncOperation(rawResponse) !== undefined ||
     getOperationLocation(rawResponse) !== undefined
   ) {
+    return {
+      mode: "Location",
+      resourceLocation: findResourceLocation(requestMethod, rawResponse, requestPath),
+    };
+  } else if (getLocation(rawResponse) !== undefined) {
     return {
       mode: "Location",
     };

sdk/core/core-lro/src/lroEngine/stateMachine.ts

@@ -14,7 +14,6 @@ import {
 import { getPollingUrl, inferLroMode, isUnexpectedInitialResponse } from "./requestUtils";
 import { isBodyPollingDone, processBodyPollingOperationResult } from "./bodyPolling";
 import { logger } from "./logger";
-import { processAzureAsyncOperationResult } from "./azureAsyncPolling";
 import { processLocationPollingOperationResult } from "./locationPolling";
 import { processPassthroughOperationResult } from "./passthrough";
 
@@ -27,16 +26,13 @@ export function createGetLroStatusFromResponse<TResult>(
   lroResourceLocationConfig?: LroResourceLocationConfig
 ): GetLroStatusFromResponse<TResult> {
   switch (config.mode) {
-    case "AzureAsync": {
-      return processAzureAsyncOperationResult(
+    case "Location": {
+      return processLocationPollingOperationResult(
         lroPrimitives,
         config.resourceLocation,
         lroResourceLocationConfig
       );
     }
-    case "Location": {
-      return processLocationPollingOperationResult;
-    }
     case "Body": {
       return processBodyPollingOperationResult;
     }