From fb160d059e69978f0279b1e79100e530285cd24d Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Tue, 22 Mar 2022 09:34:33 -0700
Subject: [PATCH] chore(middleware-retry): use normalizeProvider from
 util-middleware (#3458)

---
 packages/middleware-retry/package.json        |  1 +
 .../src/configurations.spec.ts                | 23 +++++++++++++++----
 .../middleware-retry/src/configurations.ts    | 20 +++-------------
 3 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/packages/middleware-retry/package.json b/packages/middleware-retry/package.json
index 861daf589291..433fb0346fad 100644
--- a/packages/middleware-retry/package.json
+++ b/packages/middleware-retry/package.json
@@ -21,6 +21,7 @@
   "dependencies": {
     "@aws-sdk/protocol-http": "*",
     "@aws-sdk/service-error-classification": "*",
+    "@aws-sdk/util-middleware": "*",
     "@aws-sdk/types": "*",
     "tslib": "^2.3.1",
     "uuid": "^8.3.2"
diff --git a/packages/middleware-retry/src/configurations.spec.ts b/packages/middleware-retry/src/configurations.spec.ts
index d1c04d786b9e..649e7a51033c 100644
--- a/packages/middleware-retry/src/configurations.spec.ts
+++ b/packages/middleware-retry/src/configurations.spec.ts
@@ -1,3 +1,5 @@
+import { normalizeProvider } from "@aws-sdk/util-middleware";
+
 import { AdaptiveRetryStrategy } from "./AdaptiveRetryStrategy";
 import { DEFAULT_MAX_ATTEMPTS } from "./config";
 import {
@@ -10,19 +12,30 @@ import { StandardRetryStrategy } from "./StandardRetryStrategy";
 
 jest.mock("./AdaptiveRetryStrategy");
 jest.mock("./StandardRetryStrategy");
+jest.mock("@aws-sdk/util-middleware");
 
 describe(resolveRetryConfig.name, () => {
   const retryMode = jest.fn() as any;
+
+  beforeEach(() => {
+    (normalizeProvider as jest.Mock).mockImplementation((input) =>
+      typeof input === "function" ? input : () => Promise.resolve(input)
+    );
+  });
+
   afterEach(() => {
     jest.clearAllMocks();
   });
 
   describe("maxAttempts", () => {
-    it("assigns maxAttempts value if present", async () => {
-      for (const maxAttempts of [1, 2, 3]) {
-        const output = await resolveRetryConfig({ maxAttempts, retryMode }).maxAttempts();
-        expect(output).toStrictEqual(maxAttempts);
-      }
+    it.each([1, 2, 3])("assigns provided value %s", async (maxAttempts) => {
+      const output = await resolveRetryConfig({ maxAttempts, retryMode }).maxAttempts();
+      expect(output).toStrictEqual(maxAttempts);
+    });
+
+    it(`assigns default ${DEFAULT_MAX_ATTEMPTS} is value not provided`, async () => {
+      const output = await resolveRetryConfig({ retryMode }).maxAttempts();
+      expect(output).toStrictEqual(DEFAULT_MAX_ATTEMPTS);
     });
   });
 
diff --git a/packages/middleware-retry/src/configurations.ts b/packages/middleware-retry/src/configurations.ts
index 9de4faee555b..c988594914bb 100644
--- a/packages/middleware-retry/src/configurations.ts
+++ b/packages/middleware-retry/src/configurations.ts
@@ -1,5 +1,6 @@
 import { LoadedConfigSelectors } from "@aws-sdk/node-config-provider";
 import { Provider, RetryStrategy } from "@aws-sdk/types";
+import { normalizeProvider } from "@aws-sdk/util-middleware";
 
 import { AdaptiveRetryStrategy } from "./AdaptiveRetryStrategy";
 import { DEFAULT_MAX_ATTEMPTS, DEFAULT_RETRY_MODE, RETRY_MODES } from "./config";
@@ -61,7 +62,7 @@ export interface RetryResolvedConfig {
 }
 
 export const resolveRetryConfig = <T>(input: T & PreviouslyResolved & RetryInputConfig): T & RetryResolvedConfig => {
-  const maxAttempts = normalizeMaxAttempts(input.maxAttempts);
+  const maxAttempts = normalizeProvider(input.maxAttempts ?? DEFAULT_MAX_ATTEMPTS);
   return {
     ...input,
     maxAttempts,
@@ -69,7 +70,7 @@ export const resolveRetryConfig = <T>(input: T & PreviouslyResolved & RetryInput
       if (input.retryStrategy) {
         return input.retryStrategy;
       }
-      const retryMode = await getRetryMode(input.retryMode);
+      const retryMode = await normalizeProvider(input.retryMode)();
       if (retryMode === RETRY_MODES.ADAPTIVE) {
         return new AdaptiveRetryStrategy(maxAttempts);
       }
@@ -78,21 +79,6 @@ export const resolveRetryConfig = <T>(input: T & PreviouslyResolved & RetryInput
   };
 };
 
-const getRetryMode = async (retryMode: string | Provider<string>): Promise<string> => {
-  if (typeof retryMode === "string") {
-    return retryMode;
-  }
-  return await retryMode();
-};
-
-const normalizeMaxAttempts = (maxAttempts: number | Provider<number> = DEFAULT_MAX_ATTEMPTS): Provider<number> => {
-  if (typeof maxAttempts === "number") {
-    const promisified = Promise.resolve(maxAttempts);
-    return () => promisified;
-  }
-  return maxAttempts;
-};
-
 export const ENV_RETRY_MODE = "AWS_RETRY_MODE";
 export const CONFIG_RETRY_MODE = "retry_mode";