From 3e782527d0cf77889fb2fed010daebec9bf1e737 Mon Sep 17 00:00:00 2001
From: victor-meng <56978073+victor-meng@users.noreply.github.com>
Date: Mon, 2 Nov 2020 16:59:08 -0800
Subject: [PATCH] Poll on Location header for operation status (#309)

---
 src/Utils/arm/request.test.ts | 10 +++++-----
 src/Utils/arm/request.ts      | 36 ++++++++---------------------------
 2 files changed, 13 insertions(+), 33 deletions(-)

diff --git a/src/Utils/arm/request.test.ts b/src/Utils/arm/request.test.ts
index 69f0be223..89b6d488a 100644
--- a/src/Utils/arm/request.test.ts
+++ b/src/Utils/arm/request.test.ts
@@ -21,13 +21,12 @@ describe("ARM request", () => {
 
   it("should poll for async operations", async () => {
     const headers = new Headers();
-    headers.set("azure-asyncoperation", "https://foo.com/operationStatus");
+    headers.set("location", "https://foo.com/operationStatus");
     window.fetch = jest.fn().mockResolvedValue({
       ok: true,
       headers,
-      json: async () => {
-        return { status: "Succeeded" };
-      }
+      status: 200,
+      json: async () => ({})
     });
     await armRequest({ apiVersion: "2001-01-01", host: "https://foo.com", path: "foo", method: "GET" });
     expect(window.fetch).toHaveBeenCalledTimes(2);
@@ -35,10 +34,11 @@ describe("ARM request", () => {
 
   it("should throw for failed async operations", async () => {
     const headers = new Headers();
-    headers.set("azure-asyncoperation", "https://foo.com/operationStatus");
+    headers.set("location", "https://foo.com/operationStatus");
     window.fetch = jest.fn().mockResolvedValue({
       ok: true,
       headers,
+      status: 200,
       json: async () => {
         return { status: "Failed" };
       }
diff --git a/src/Utils/arm/request.ts b/src/Utils/arm/request.ts
index d41803e37..e40e38b8f 100644
--- a/src/Utils/arm/request.ts
+++ b/src/Utils/arm/request.ts
@@ -70,55 +70,35 @@ export async function armRequest<T>({ host, path, apiVersion, method, body: requ
     throw error;
   }
 
-  const operationStatusUrl = response.headers && response.headers.get("azure-asyncoperation");
+  const operationStatusUrl = response.headers && response.headers.get("location");
   if (operationStatusUrl) {
-    await promiseRetry(() => getOperationStatus(operationStatusUrl));
-    // TODO: ARM is supposed to return a resourceLocation property, but it does not https://github.com/microsoft/api-guidelines/blob/vNext/Guidelines.md#target-resource-location
-    // When Cosmos RP adds resourceLocation, we should use it instead
-    // For now manually execute a GET if the operation was a mutation and not a deletion
-    if (method === "POST" || method === "PATCH" || method === "PUT") {
-      return armRequest({
-        host,
-        path,
-        apiVersion,
-        method: "GET"
-      });
-    }
+    return await promiseRetry(() => getOperationStatus(operationStatusUrl));
   }
 
   const responseBody = (await response.json()) as T;
   return responseBody;
 }
 
-const SUCCEEDED = "Succeeded" as const;
-const FAILED = "Failed" as const;
-const CANCELED = "Canceled" as const;
-
-type Status = typeof SUCCEEDED | typeof FAILED | typeof CANCELED;
-
-interface OperationResponse {
-  status: Status;
-  error: unknown;
-}
-
 async function getOperationStatus(operationStatusUrl: string) {
   const response = await window.fetch(operationStatusUrl, {
     headers: {
       Authorization: userContext.authorizationToken
     }
   });
+
   if (!response.ok) {
     const errorResponse = (await response.json()) as ErrorResponse;
     const error = new Error(errorResponse.message) as ARMError;
     error.code = errorResponse.code;
     throw new AbortError(error);
   }
-  const body = (await response.json()) as OperationResponse;
+
+  const body = await response.json();
   const status = body.status;
-  if (status === SUCCEEDED) {
-    return;
+  if (!status && response.status === 200) {
+    return body;
   }
-  if (status === CANCELED || status === FAILED) {
+  if (status === "Canceled" || status === "Failed") {
     const errorMessage = body.error ? JSON.stringify(body.error) : "Operation could not be completed";
     const error = new Error(errorMessage);
     throw new AbortError(error);