From a021d43d14949dcf58307192303e544d634c49b7 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Wed, 10 Jun 2020 11:01:45 -0700
Subject: [PATCH 1/7] test(middleware-retry): add extensive unit tests (#1255)

---
 .../src/configurations.spec.ts                |  47 ++++
 .../middleware-retry/src/configurations.ts    |   8 +-
 .../src/defaultStrategy.spec.ts               | 203 ++++++++++++++++++
 packages/middleware-retry/src/index.spec.ts   |  92 --------
 .../middleware-retry/src/retryDecider.spec.ts |  48 ++---
 .../src/retryMiddleware.spec.ts               |  76 +++++++
 .../middleware-retry/src/retryMiddleware.ts   |  17 +-
 7 files changed, 362 insertions(+), 129 deletions(-)
 create mode 100644 packages/middleware-retry/src/configurations.spec.ts
 create mode 100644 packages/middleware-retry/src/defaultStrategy.spec.ts
 delete mode 100644 packages/middleware-retry/src/index.spec.ts
 create mode 100644 packages/middleware-retry/src/retryMiddleware.spec.ts

diff --git a/packages/middleware-retry/src/configurations.spec.ts b/packages/middleware-retry/src/configurations.spec.ts
new file mode 100644
index 000000000000..2559251b6483
--- /dev/null
+++ b/packages/middleware-retry/src/configurations.spec.ts
@@ -0,0 +1,47 @@
+import { resolveRetryConfig } from "./configurations";
+import { StandardRetryStrategy } from "./defaultStrategy";
+
+describe("resolveRetryConfig", () => {
+  describe("maxAttempts", () => {
+    it("uses passed maxAttempts value if present", () => {
+      [1, 2, 3].forEach(maxAttempts => {
+        expect(resolveRetryConfig({ maxAttempts }).maxAttempts).toEqual(
+          maxAttempts
+        );
+      });
+    });
+
+    it("assigns default value of 3 if maxAttempts not passed", () => {
+      expect(resolveRetryConfig({}).maxAttempts).toEqual(3);
+    });
+  });
+
+  describe("retryStrategy", () => {
+    it("uses passed retryStrategy if present", () => {
+      const mockRetryStrategy = {
+        maxAttempts: 2,
+        retry: jest.fn()
+      };
+      const { retryStrategy } = resolveRetryConfig({
+        retryStrategy: mockRetryStrategy
+      });
+      expect(retryStrategy).toEqual(mockRetryStrategy);
+    });
+
+    describe("creates StandardRetryStrategy if retryStrategy not present", () => {
+      describe("uses maxAttempts if present", () => {
+        [1, 2, 3].forEach(maxAttempts => {
+          const { retryStrategy } = resolveRetryConfig({ maxAttempts });
+          expect(retryStrategy).toBeInstanceOf(StandardRetryStrategy);
+          expect(retryStrategy.maxAttempts).toBe(maxAttempts);
+        });
+      });
+
+      it("uses default 3 if maxAttempts is not present", () => {
+        const { retryStrategy } = resolveRetryConfig({});
+        expect(retryStrategy).toBeInstanceOf(StandardRetryStrategy);
+        expect(retryStrategy.maxAttempts).toBe(3);
+      });
+    });
+  });
+});
diff --git a/packages/middleware-retry/src/configurations.ts b/packages/middleware-retry/src/configurations.ts
index 8a5eb0178551..a52a18b83420 100644
--- a/packages/middleware-retry/src/configurations.ts
+++ b/packages/middleware-retry/src/configurations.ts
@@ -17,13 +17,13 @@ export interface RetryResolvedConfig {
   retryStrategy: RetryStrategy;
 }
 
-export function resolveRetryConfig<T>(
+export const resolveRetryConfig = <T>(
   input: T & RetryInputConfig
-): T & RetryResolvedConfig {
-  const maxAttempts = input.maxAttempts === undefined ? 3 : input.maxAttempts;
+): T & RetryResolvedConfig => {
+  const maxAttempts = input.maxAttempts ?? 3;
   return {
     ...input,
     maxAttempts,
     retryStrategy: input.retryStrategy || new StandardRetryStrategy(maxAttempts)
   };
-}
+};
diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/defaultStrategy.spec.ts
new file mode 100644
index 000000000000..0b9207a34eeb
--- /dev/null
+++ b/packages/middleware-retry/src/defaultStrategy.spec.ts
@@ -0,0 +1,203 @@
+import {
+  DEFAULT_RETRY_DELAY_BASE,
+  THROTTLING_RETRY_DELAY_BASE
+} from "./constants";
+import { isThrottlingError } from "@aws-sdk/service-error-classification";
+import { defaultDelayDecider } from "./delayDecider";
+import { defaultRetryDecider } from "./retryDecider";
+import { StandardRetryStrategy } from "./defaultStrategy";
+
+jest.mock("@aws-sdk/service-error-classification", () => ({
+  isThrottlingError: jest.fn().mockReturnValue(true)
+}));
+
+jest.mock("./delayDecider", () => ({
+  defaultDelayDecider: jest.fn().mockReturnValue(0)
+}));
+
+jest.mock("./retryDecider", () => ({
+  defaultRetryDecider: jest.fn().mockReturnValue(true)
+}));
+
+describe("defaultStrategy", () => {
+  const maxAttempts = 2;
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it("sets maxAttempts as class member variable", () => {
+    [1, 2, 3].forEach(maxAttempts => {
+      const retryStrategy = new StandardRetryStrategy(maxAttempts);
+      expect(retryStrategy.maxAttempts).toBe(maxAttempts);
+    });
+  });
+
+  describe("retryDecider", () => {
+    it("sets defaultRetryDecider if options is undefined", () => {
+      const retryStrategy = new StandardRetryStrategy(maxAttempts);
+      expect(retryStrategy["retryDecider"]).toBe(defaultRetryDecider);
+    });
+
+    it("sets defaultRetryDecider if options.retryDecider is undefined", () => {
+      const retryStrategy = new StandardRetryStrategy(maxAttempts, {});
+      expect(retryStrategy["retryDecider"]).toBe(defaultRetryDecider);
+    });
+
+    it("sets options.retryDecider if defined", () => {
+      const retryDecider = jest.fn();
+      const retryStrategy = new StandardRetryStrategy(maxAttempts, {
+        retryDecider
+      });
+      expect(retryStrategy["retryDecider"]).toBe(retryDecider);
+    });
+  });
+
+  describe("delayDecider", () => {
+    it("sets defaultDelayDecider if options is undefined", () => {
+      const retryStrategy = new StandardRetryStrategy(maxAttempts);
+      expect(retryStrategy["delayDecider"]).toBe(defaultDelayDecider);
+    });
+
+    it("sets defaultDelayDecider if options.delayDecider undefined", () => {
+      const retryStrategy = new StandardRetryStrategy(maxAttempts, {});
+      expect(retryStrategy["delayDecider"]).toBe(defaultDelayDecider);
+    });
+
+    it("sets options.delayDecider if defined", () => {
+      const delayDecider = jest.fn();
+      const retryStrategy = new StandardRetryStrategy(maxAttempts, {
+        delayDecider
+      });
+      expect(retryStrategy["delayDecider"]).toBe(delayDecider);
+    });
+  });
+
+  describe("delayBase passed to delayDecider", () => {
+    const testDelayBasePassed = async (
+      delayBaseToTest: number,
+      mockThrottlingError: boolean
+    ) => {
+      (isThrottlingError as jest.Mock).mockReturnValueOnce(mockThrottlingError);
+
+      const mockError = new Error("mockError");
+      const mockResponse = {
+        response: "mockResponse",
+        output: { $metadata: {} }
+      };
+
+      const next = jest
+        .fn()
+        .mockRejectedValueOnce(mockError)
+        .mockResolvedValueOnce(mockResponse);
+
+      const retryStrategy = new StandardRetryStrategy(maxAttempts);
+      await retryStrategy.retry(next, {} as any);
+
+      expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1);
+      expect(isThrottlingError as jest.Mock).toHaveBeenCalledWith(mockError);
+      expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(1);
+      expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledWith(
+        delayBaseToTest,
+        1
+      );
+    };
+
+    it("should be equal to THROTTLING_RETRY_DELAY_BASE if error is throttling error", async () => {
+      return testDelayBasePassed(THROTTLING_RETRY_DELAY_BASE, true);
+    });
+
+    it("should be equal to DEFAULT_RETRY_DELAY_BASE in error is not a throttling error", async () => {
+      return testDelayBasePassed(DEFAULT_RETRY_DELAY_BASE, false);
+    });
+  });
+
+  describe("should not retry", () => {
+    it("when the handler completes successfully", async () => {
+      const mockResponse = {
+        response: "mockResponse",
+        output: { $metadata: {} }
+      };
+
+      const next = jest.fn().mockResolvedValueOnce(mockResponse);
+
+      const retryStrategy = new StandardRetryStrategy(maxAttempts);
+      const { response, output } = await retryStrategy.retry(next, {} as any);
+
+      expect(response).toStrictEqual(mockResponse.response);
+      expect(output.$metadata.attempts).toBe(1);
+      expect(output.$metadata.totalRetryDelay).toBe(0);
+      expect(defaultRetryDecider as jest.Mock).not.toHaveBeenCalled();
+      expect(defaultDelayDecider as jest.Mock).not.toHaveBeenCalled();
+    });
+
+    it("when retryDecider returns false", async () => {
+      (defaultRetryDecider as jest.Mock).mockReturnValueOnce(false);
+
+      const mockError = new Error("mockError");
+      const next = jest.fn().mockRejectedValueOnce(mockError);
+
+      const retryStrategy = new StandardRetryStrategy(maxAttempts);
+      try {
+        await retryStrategy.retry(next, {} as any);
+      } catch (error) {
+        expect(error).toStrictEqual(mockError);
+      }
+
+      expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledTimes(1);
+      expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledWith(mockError);
+    });
+
+    it("when the the maximum number of attempts is reached", async () => {
+      const mockError = new Error("mockError");
+      const next = jest.fn().mockRejectedValue(mockError);
+
+      const retryStrategy = new StandardRetryStrategy(maxAttempts);
+      try {
+        await retryStrategy.retry(next, {} as any);
+      } catch (error) {
+        expect(error).toStrictEqual(mockError);
+      }
+      expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledTimes(
+        maxAttempts - 1
+      );
+    });
+  });
+
+  it("should delay equal to the value returned by delayDecider", async () => {
+    jest.spyOn(global, "setTimeout");
+
+    const FIRST_DELAY = 100;
+    const SECOND_DELAY = 200;
+
+    (defaultDelayDecider as jest.Mock)
+      .mockReturnValueOnce(FIRST_DELAY)
+      .mockReturnValueOnce(SECOND_DELAY);
+
+    const mockError = new Error("mockError");
+    const next = jest.fn().mockRejectedValue(mockError);
+
+    const retryStrategy = new StandardRetryStrategy(3);
+    try {
+      await retryStrategy.retry(next, {} as any);
+    } catch (error) {
+      expect(error).toStrictEqual(mockError);
+      expect(error.$metadata.totalRetryDelay).toEqual(
+        FIRST_DELAY + SECOND_DELAY
+      );
+    }
+
+    expect(defaultDelayDecider as jest.Mock).toHaveBeenCalledTimes(2);
+    expect(setTimeout).toHaveBeenCalledTimes(2);
+    expect(setTimeout).toHaveBeenNthCalledWith(
+      1,
+      expect.any(Function),
+      FIRST_DELAY
+    );
+    expect(setTimeout).toHaveBeenNthCalledWith(
+      2,
+      expect.any(Function),
+      SECOND_DELAY
+    );
+  });
+});
diff --git a/packages/middleware-retry/src/index.spec.ts b/packages/middleware-retry/src/index.spec.ts
deleted file mode 100644
index 3c1f972b120a..000000000000
--- a/packages/middleware-retry/src/index.spec.ts
+++ /dev/null
@@ -1,92 +0,0 @@
-import {
-  DEFAULT_RETRY_DELAY_BASE,
-  THROTTLING_RETRY_DELAY_BASE
-} from "./constants";
-import { retryMiddleware } from "./retryMiddleware";
-import { resolveRetryConfig } from "./configurations";
-import * as delayDeciderModule from "./delayDecider";
-import { StandardRetryStrategy, RetryDecider } from "./defaultStrategy";
-import { HttpRequest } from "@aws-sdk/protocol-http";
-import { SdkError } from "@aws-sdk/smithy-client";
-
-describe("retryMiddleware", () => {
-  it("should not retry when the handler completes successfully", async () => {
-    const next = jest.fn().mockResolvedValue({ output: { $metadata: {} } });
-    const retryHandler = retryMiddleware(
-      resolveRetryConfig({ maxAttempts: 0 })
-    )(next);
-
-    const {
-      output: { $metadata }
-    } = await retryHandler({ input: {}, request: new HttpRequest({}) });
-    expect($metadata.attempts).toBe(1);
-    expect($metadata.totalRetryDelay).toBe(0);
-
-    expect(next.mock.calls.length).toBe(1);
-  });
-
-  it("should stop retrying when the the maximum number of retries is reached", async () => {
-    const maxAttempts = 3;
-    const error = new Error();
-    error.name = "ProvisionedThroughputExceededException";
-    const next = jest.fn().mockRejectedValue(error);
-    const retryHandler = retryMiddleware(resolveRetryConfig({ maxAttempts }))(
-      next
-    );
-
-    await expect(
-      retryHandler({ input: {}, request: new HttpRequest({}) })
-    ).rejects.toMatchObject(error);
-
-    expect(next.mock.calls.length).toBe(maxAttempts);
-  });
-
-  it("should not retry if the error is not transient", async () => {
-    const error = new Error();
-    error.name = "ValidationException";
-    const next = jest.fn().mockRejectedValue(error);
-    const retryHandler = retryMiddleware(
-      resolveRetryConfig({ maxAttempts: 3 })
-    )(next);
-
-    await expect(
-      retryHandler({ input: {}, request: new HttpRequest({}) })
-    ).rejects.toMatchObject(error);
-
-    expect(next.mock.calls.length).toBe(1);
-  });
-
-  it("should use a higher base delay when a throttling error is encountered", async () => {
-    const next = jest.fn().mockResolvedValue({ output: { $metadata: {} } });
-
-    const validation = new Error();
-    validation.name = "ValidationException";
-    next.mockImplementationOnce(args => Promise.reject(validation));
-
-    const throttling = new Error();
-    throttling.name = "RequestLimitExceeded";
-    next.mockImplementationOnce(args => Promise.reject(throttling));
-
-    jest.mock("./delayDecider");
-
-    const maxAttempts = 3;
-    const delayDeciderMock = jest.spyOn(
-      delayDeciderModule,
-      "defaultDelayDecider"
-    );
-    const retryDecider: RetryDecider = (error: SdkError) => true;
-    const strategy = new StandardRetryStrategy(maxAttempts, { retryDecider });
-    const retryHandler = retryMiddleware({
-      maxAttempts,
-      retryStrategy: strategy
-    })(next);
-
-    await retryHandler({ input: {}, request: new HttpRequest({}) });
-
-    expect(next.mock.calls.length).toBe(3);
-    expect(delayDeciderMock.mock.calls).toEqual([
-      [DEFAULT_RETRY_DELAY_BASE, 1],
-      [THROTTLING_RETRY_DELAY_BASE, 2]
-    ]);
-  });
-});
diff --git a/packages/middleware-retry/src/retryDecider.spec.ts b/packages/middleware-retry/src/retryDecider.spec.ts
index 15ed0732cfee..271f22baba0b 100644
--- a/packages/middleware-retry/src/retryDecider.spec.ts
+++ b/packages/middleware-retry/src/retryDecider.spec.ts
@@ -24,53 +24,53 @@ describe("defaultRetryDecider", () => {
 
   it("should return false when the provided error is falsy", () => {
     expect(defaultRetryDecider(null as any)).toBe(false);
-    expect((isRetryableByTrait as jest.Mock).mock.calls.length).toBe(0);
-    expect((isClockSkewError as jest.Mock).mock.calls.length).toBe(0);
-    expect((isThrottlingError as jest.Mock).mock.calls.length).toBe(0);
-    expect((isTransientError as jest.Mock).mock.calls.length).toBe(0);
+    expect(isRetryableByTrait as jest.Mock).toHaveBeenCalledTimes(0);
+    expect(isClockSkewError as jest.Mock).toHaveBeenCalledTimes(0);
+    expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(0);
+    expect(isTransientError as jest.Mock).toHaveBeenCalledTimes(0);
   });
 
   it("should return true for RetryableByTrait error", () => {
     (isRetryableByTrait as jest.Mock).mockReturnValueOnce(true);
     expect(defaultRetryDecider(createMockError())).toBe(true);
-    expect((isRetryableByTrait as jest.Mock).mock.calls.length).toBe(1);
-    expect((isClockSkewError as jest.Mock).mock.calls.length).toBe(0);
-    expect((isThrottlingError as jest.Mock).mock.calls.length).toBe(0);
-    expect((isTransientError as jest.Mock).mock.calls.length).toBe(0);
+    expect(isRetryableByTrait as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isClockSkewError as jest.Mock).toHaveBeenCalledTimes(0);
+    expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(0);
+    expect(isTransientError as jest.Mock).toHaveBeenCalledTimes(0);
   });
 
   it("should return true for ClockSkewError", () => {
     (isClockSkewError as jest.Mock).mockReturnValueOnce(true);
     expect(defaultRetryDecider(createMockError())).toBe(true);
-    expect((isRetryableByTrait as jest.Mock).mock.calls.length).toBe(1);
-    expect((isClockSkewError as jest.Mock).mock.calls.length).toBe(1);
-    expect((isThrottlingError as jest.Mock).mock.calls.length).toBe(0);
-    expect((isTransientError as jest.Mock).mock.calls.length).toBe(0);
+    expect(isRetryableByTrait as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isClockSkewError as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(0);
+    expect(isTransientError as jest.Mock).toHaveBeenCalledTimes(0);
   });
 
   it("should return true for ThrottlingError", () => {
     (isThrottlingError as jest.Mock).mockReturnValueOnce(true);
     expect(defaultRetryDecider(createMockError())).toBe(true);
-    expect((isRetryableByTrait as jest.Mock).mock.calls.length).toBe(1);
-    expect((isClockSkewError as jest.Mock).mock.calls.length).toBe(1);
-    expect((isThrottlingError as jest.Mock).mock.calls.length).toBe(1);
-    expect((isTransientError as jest.Mock).mock.calls.length).toBe(0);
+    expect(isRetryableByTrait as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isClockSkewError as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isTransientError as jest.Mock).toHaveBeenCalledTimes(0);
   });
 
   it("should return true for TransientError", () => {
     (isTransientError as jest.Mock).mockReturnValueOnce(true);
     expect(defaultRetryDecider(createMockError())).toBe(true);
-    expect((isRetryableByTrait as jest.Mock).mock.calls.length).toBe(1);
-    expect((isClockSkewError as jest.Mock).mock.calls.length).toBe(1);
-    expect((isThrottlingError as jest.Mock).mock.calls.length).toBe(1);
-    expect((isTransientError as jest.Mock).mock.calls.length).toBe(1);
+    expect(isRetryableByTrait as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isClockSkewError as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isTransientError as jest.Mock).toHaveBeenCalledTimes(1);
   });
 
   it("should return false for other errors", () => {
     expect(defaultRetryDecider(createMockError())).toBe(false);
-    expect((isRetryableByTrait as jest.Mock).mock.calls.length).toBe(1);
-    expect((isClockSkewError as jest.Mock).mock.calls.length).toBe(1);
-    expect((isThrottlingError as jest.Mock).mock.calls.length).toBe(1);
-    expect((isTransientError as jest.Mock).mock.calls.length).toBe(1);
+    expect(isRetryableByTrait as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isClockSkewError as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1);
+    expect(isTransientError as jest.Mock).toHaveBeenCalledTimes(1);
   });
 });
diff --git a/packages/middleware-retry/src/retryMiddleware.spec.ts b/packages/middleware-retry/src/retryMiddleware.spec.ts
new file mode 100644
index 000000000000..08a49b3708e5
--- /dev/null
+++ b/packages/middleware-retry/src/retryMiddleware.spec.ts
@@ -0,0 +1,76 @@
+import {
+  getRetryPlugin,
+  retryMiddleware,
+  retryMiddlewareOptions
+} from "./retryMiddleware";
+import {
+  MiddlewareStack,
+  RetryStrategy,
+  FinalizeHandlerArguments
+} from "@aws-sdk/types";
+
+describe("getRetryPlugin", () => {
+  const mockClientStack = {
+    add: jest.fn()
+  };
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe("adds retryMiddleware if maxAttempts > 1", () => {
+    [2, 3, 4].forEach(maxAttempts => {
+      it(`when maxAttempts=${maxAttempts}`, () => {
+        getRetryPlugin({
+          maxAttempts,
+          retryStrategy: {} as RetryStrategy
+        }).applyToStack(
+          (mockClientStack as unknown) as MiddlewareStack<any, any>
+        );
+        expect(mockClientStack.add).toHaveBeenCalledTimes(1);
+        expect(mockClientStack.add.mock.calls[0][1]).toEqual(
+          retryMiddlewareOptions
+        );
+      });
+    });
+  });
+
+  describe("skips adding retryMiddleware if maxAttempts <= 1", () => {
+    [0, 1].forEach(maxAttempts => {
+      it(`when maxAttempts=${maxAttempts}`, () => {
+        getRetryPlugin({
+          maxAttempts,
+          retryStrategy: {} as RetryStrategy
+        }).applyToStack(
+          (mockClientStack as unknown) as MiddlewareStack<any, any>
+        );
+        expect(mockClientStack.add).toHaveBeenCalledTimes(0);
+      });
+    });
+  });
+});
+
+describe("retryMiddleware", () => {
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it("calls retryStrategy.retry with next and args", async () => {
+    const maxAttempts = 2;
+    const next = jest.fn();
+    const args = {
+      request: {}
+    };
+    const mockRetryStrategy = {
+      maxAttempts,
+      retry: jest.fn()
+    };
+
+    await retryMiddleware({
+      maxAttempts,
+      retryStrategy: mockRetryStrategy
+    })(next)(args as FinalizeHandlerArguments<any>);
+    expect(mockRetryStrategy.retry).toHaveBeenCalledTimes(1);
+    expect(mockRetryStrategy.retry).toHaveBeenCalledWith(next, args);
+  });
+});
diff --git a/packages/middleware-retry/src/retryMiddleware.ts b/packages/middleware-retry/src/retryMiddleware.ts
index 2bc454e80506..bbd72ec957f4 100644
--- a/packages/middleware-retry/src/retryMiddleware.ts
+++ b/packages/middleware-retry/src/retryMiddleware.ts
@@ -9,15 +9,14 @@ import {
 } from "@aws-sdk/types";
 import { RetryResolvedConfig } from "./configurations";
 
-export function retryMiddleware(options: RetryResolvedConfig) {
-  return <Output extends MetadataBearer = MetadataBearer>(
-    next: FinalizeHandler<any, Output>
-  ): FinalizeHandler<any, Output> => async (
-    args: FinalizeHandlerArguments<any>
-  ): Promise<FinalizeHandlerOutput<Output>> => {
-    return options.retryStrategy.retry(next, args);
-  };
-}
+export const retryMiddleware = (options: RetryResolvedConfig) => <
+  Output extends MetadataBearer = MetadataBearer
+>(
+  next: FinalizeHandler<any, Output>
+): FinalizeHandler<any, Output> => async (
+  args: FinalizeHandlerArguments<any>
+): Promise<FinalizeHandlerOutput<Output>> =>
+  options.retryStrategy.retry(next, args);
 
 export const retryMiddlewareOptions: FinalizeRequestHandlerOptions &
   AbsoluteLocation = {

From cd4813013063fc015678bfb1e71d7c1bbc26e2ab Mon Sep 17 00:00:00 2001
From: "dependabot-preview[bot]"
 <27856297+dependabot-preview[bot]@users.noreply.github.com>
Date: Wed, 10 Jun 2020 18:05:44 +0000
Subject: [PATCH 2/7] chore(deps-dev): bump prettier from 2.0.4 to 2.0.5
 (#1200)

---
 package.json | 2 +-
 yarn.lock    | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/package.json b/package.json
index 97d5205b7119..00c6a315e9f4 100644
--- a/package.json
+++ b/package.json
@@ -49,7 +49,7 @@
     "jmespath": "^0.15.0",
     "lerna": "3.20.2",
     "lint-staged": "^10.0.1",
-    "prettier": "2.0.4",
+    "prettier": "2.0.5",
     "typescript": "~3.8.3",
     "yarn": "1.22.4"
   },
diff --git a/yarn.lock b/yarn.lock
index 89dd66f8822b..3efe1077d665 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -8176,10 +8176,10 @@ prelude-ls@~1.1.2:
   resolved "https://registry.yarnpkg.com/prelude-ls/-/prelude-ls-1.1.2.tgz#21932a549f5e52ffd9a827f570e04be62a97da54"
   integrity sha1-IZMqVJ9eUv/ZqCf1cOBL5iqX2lQ=
 
-prettier@2.0.4:
-  version "2.0.4"
-  resolved "https://registry.yarnpkg.com/prettier/-/prettier-2.0.4.tgz#2d1bae173e355996ee355ec9830a7a1ee05457ef"
-  integrity sha512-SVJIQ51spzFDvh4fIbCLvciiDMCrRhlN3mbZvv/+ycjvmF5E73bKdGfU8QDLNmjYJf+lsGnDBC4UUnvTe5OO0w==
+prettier@2.0.5:
+  version "2.0.5"
+  resolved "https://registry.yarnpkg.com/prettier/-/prettier-2.0.5.tgz#d6d56282455243f2f92cc1716692c08aa31522d4"
+  integrity sha512-7PtVymN48hGcO4fGjybyBSIWDsLU4H4XlvOHfq91pz9kkGlonzwTfYkaIEwiRg/dAJF9YlbsduBAgtYLi+8cFg==
 
 pretty-format@^24.9.0:
   version "24.9.0"

From 1155483dd0371b0997b78a401e7edd299538ee12 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Tue, 9 Jun 2020 18:31:58 +0000
Subject: [PATCH 3/7] chore: add optional RetryQuota to StandardRetryStrategy

---
 .../middleware-retry/src/defaultStrategy.ts   | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/packages/middleware-retry/src/defaultStrategy.ts b/packages/middleware-retry/src/defaultStrategy.ts
index c2b5d798ed8c..7b91e9563589 100644
--- a/packages/middleware-retry/src/defaultStrategy.ts
+++ b/packages/middleware-retry/src/defaultStrategy.ts
@@ -33,17 +33,40 @@ export interface DelayDecider {
   (delayBase: number, attempts: number): number;
 }
 
+/**
+ * Interface that specifies the retry quota behavior.
+ */
+export interface RetryQuota {
+  /**
+   * returns true if retry tokens are available from the retry quota bucket.
+   */
+  hasRetryTokens: (error: SdkError) => boolean;
+
+  /**
+   * returns token amount from the retry quota bucket.
+   * throws error is retry tokens are not available.
+   */
+  retrieveRetryToken: (error: SdkError) => number;
+
+  /**
+   * releases tokens back to the retry quota.
+   */
+  releaseRetryToken: (releaseCapacityAmount?: number) => void;
+}
+
 /**
  * Strategy options to be passed to StandardRetryStrategy
  */
 export interface StandardRetryStrategyOptions {
   retryDecider?: RetryDecider;
   delayDecider?: DelayDecider;
+  retryQuota?: RetryQuota;
 }
 
 export class StandardRetryStrategy implements RetryStrategy {
   private retryDecider: RetryDecider;
   private delayDecider: DelayDecider;
+  private retryQuota?: RetryQuota;
 
   constructor(
     public readonly maxAttempts: number,
@@ -51,6 +74,7 @@ export class StandardRetryStrategy implements RetryStrategy {
   ) {
     this.retryDecider = options?.retryDecider ?? defaultRetryDecider;
     this.delayDecider = options?.delayDecider ?? defaultDelayDecider;
+    this.retryQuota = options?.retryQuota;
   }
 
   private shouldRetry(error: SdkError, attempts: number) {

From f1064c51ea015a5bd1e6107e9673ec8c27cb1abf Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Tue, 9 Jun 2020 21:07:48 +0000
Subject: [PATCH 4/7] chore: add defaultRetryQuota

---
 packages/middleware-retry/src/constants.ts    |  22 +++
 .../src/defaultRetryQuota.spec.ts             | 147 ++++++++++++++++++
 .../middleware-retry/src/defaultRetryQuota.ts |  40 +++++
 .../middleware-retry/src/defaultStrategy.ts   |   4 +-
 4 files changed, 211 insertions(+), 2 deletions(-)
 create mode 100644 packages/middleware-retry/src/defaultRetryQuota.spec.ts
 create mode 100644 packages/middleware-retry/src/defaultRetryQuota.ts

diff --git a/packages/middleware-retry/src/constants.ts b/packages/middleware-retry/src/constants.ts
index 5f976928e514..c14baa7b87c3 100644
--- a/packages/middleware-retry/src/constants.ts
+++ b/packages/middleware-retry/src/constants.ts
@@ -15,3 +15,25 @@ export const MAXIMUM_RETRY_DELAY = 20 * 1000;
  * encountered.
  */
 export const THROTTLING_RETRY_DELAY_BASE = 500;
+
+/**
+ * Initial number of retry tokens in Retry Quota
+ */
+export const INITIAL_RETRY_TOKENS = 500;
+
+/**
+ * The total amount of retry tokens to be decremented from retry token balance.
+ */
+export const RETRY_COST = 5;
+
+/**
+ * The total amount of retry tokens to be decremented from retry token balance
+ * when a throttling error is encountered.
+ */
+export const TIMEOUT_RETRY_COST = 10;
+
+/**
+ * The total amount of retry token to be incremented from retry token balance
+ * if an SDK operation invocation succeeds without requiring a retry request.
+ */
+export const NO_RETRY_INCREMENT = 1;
diff --git a/packages/middleware-retry/src/defaultRetryQuota.spec.ts b/packages/middleware-retry/src/defaultRetryQuota.spec.ts
new file mode 100644
index 000000000000..eb6e8df7d639
--- /dev/null
+++ b/packages/middleware-retry/src/defaultRetryQuota.spec.ts
@@ -0,0 +1,147 @@
+import { defaultRetryQuota } from "./defaultRetryQuota";
+import { SdkError } from "@aws-sdk/smithy-client";
+import {
+  INITIAL_RETRY_TOKENS,
+  TIMEOUT_RETRY_COST,
+  RETRY_COST,
+  NO_RETRY_INCREMENT
+} from "./constants";
+
+describe("defaultRetryQuota", () => {
+  const getMockError = () => new Error() as SdkError;
+  const getMockTimeoutError = () =>
+    Object.assign(new Error(), {
+      name: "TimeoutError"
+    }) as SdkError;
+
+  const getDrainedRetryQuota = (targetCapacity: number, error: SdkError) => {
+    const retryQuota = defaultRetryQuota();
+    let availableCapacity = INITIAL_RETRY_TOKENS;
+    while (availableCapacity >= targetCapacity) {
+      retryQuota.retrieveRetryToken(error);
+      availableCapacity -= targetCapacity;
+    }
+    return retryQuota;
+  };
+
+  describe("hasRetryTokens", () => {
+    describe("returns true if capacity is available", () => {
+      it("when it's TimeoutError", () => {
+        const timeoutError = getMockTimeoutError();
+        expect(defaultRetryQuota().hasRetryTokens(timeoutError)).toBe(true);
+      });
+
+      it("when it's not TimeoutError", () => {
+        expect(defaultRetryQuota().hasRetryTokens(getMockError())).toBe(true);
+      });
+    });
+
+    describe("returns false if capacity is not available", () => {
+      it("when it's TimeoutError", () => {
+        const timeoutError = getMockTimeoutError();
+        const retryQuota = getDrainedRetryQuota(
+          TIMEOUT_RETRY_COST,
+          timeoutError
+        );
+        expect(retryQuota.hasRetryTokens(timeoutError)).toBe(false);
+      });
+
+      it("when it's not TimeoutError", () => {
+        const error = getMockError();
+        const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
+        expect(retryQuota.hasRetryTokens(error)).toBe(false);
+      });
+    });
+  });
+
+  describe("retrieveRetryToken", () => {
+    describe("returns retry tokens amount if available", () => {
+      it("when it's TimeoutError", () => {
+        const timeoutError = getMockTimeoutError();
+        expect(defaultRetryQuota().retrieveRetryToken(timeoutError)).toBe(
+          TIMEOUT_RETRY_COST
+        );
+      });
+
+      it("when it's not TimeoutError", () => {
+        expect(defaultRetryQuota().retrieveRetryToken(getMockError())).toBe(
+          RETRY_COST
+        );
+      });
+    });
+
+    describe("throws error if retry tokens not available", () => {
+      it("when it's TimeoutError", () => {
+        const timeoutError = getMockTimeoutError();
+        const retryQuota = getDrainedRetryQuota(
+          TIMEOUT_RETRY_COST,
+          timeoutError
+        );
+        expect(() => {
+          retryQuota.retrieveRetryToken(timeoutError);
+        }).toThrowError(new Error("No retry token available"));
+      });
+
+      it("when it's not TimeoutError", () => {
+        const error = getMockError();
+        const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
+        expect(() => {
+          retryQuota.retrieveRetryToken(error);
+        }).toThrowError(new Error("No retry token available"));
+      });
+    });
+  });
+
+  describe("releaseRetryToken", () => {
+    it("adds capacityReleaseAmount if passed", () => {
+      const error = getMockError();
+      const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
+
+      // Ensure that retry tokens are not available.
+      expect(retryQuota.hasRetryTokens(error)).toBe(false);
+
+      // Release RETRY_COST tokens.
+      retryQuota.releaseRetryToken(RETRY_COST);
+      expect(retryQuota.hasRetryTokens(error)).toBe(true);
+      expect(retryQuota.retrieveRetryToken(error)).toBe(RETRY_COST);
+      expect(retryQuota.hasRetryTokens(error)).toBe(false);
+    });
+
+    it("adds NO_RETRY_INCREMENT if capacityReleaseAmount not passed", () => {
+      const error = getMockError();
+      const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
+
+      // retry tokens will not be available till NO_RETRY_INCREMENT is added
+      // till it's equal to RETRY_COST - (INITIAL_RETRY_TOKENS % RETRY_COST)
+      let tokensReleased = 0;
+      const tokensToBeReleased =
+        RETRY_COST - (INITIAL_RETRY_TOKENS % RETRY_COST);
+      while (tokensReleased < tokensToBeReleased) {
+        expect(retryQuota.hasRetryTokens(error)).toBe(false);
+        retryQuota.releaseRetryToken();
+        tokensReleased += NO_RETRY_INCREMENT;
+      }
+      expect(retryQuota.hasRetryTokens(error)).toBe(true);
+    });
+
+    it("ensures availableCapacity is maxed at INITIAL_RETRY_TOKENS", () => {
+      const error = getMockError();
+      const retryQuota = defaultRetryQuota();
+
+      // release 100 tokens.
+      [...Array(100).keys()].forEach(key => {
+        retryQuota.releaseRetryToken();
+      });
+
+      // availableCapacity is still maxed at INITIAL_RETRY_TOKENS
+      // hasRetryTokens would be true only till INITIAL_RETRY_TOKENS/RETRY_COST times
+      [...Array(Math.floor(INITIAL_RETRY_TOKENS / RETRY_COST)).keys()].forEach(
+        key => {
+          expect(retryQuota.hasRetryTokens(error)).toBe(true);
+          retryQuota.retrieveRetryToken(error);
+        }
+      );
+      expect(retryQuota.hasRetryTokens(error)).toBe(false);
+    });
+  });
+});
diff --git a/packages/middleware-retry/src/defaultRetryQuota.ts b/packages/middleware-retry/src/defaultRetryQuota.ts
new file mode 100644
index 000000000000..a2044b02ffdd
--- /dev/null
+++ b/packages/middleware-retry/src/defaultRetryQuota.ts
@@ -0,0 +1,40 @@
+import { RetryQuota } from "./defaultStrategy";
+import { SdkError } from "@aws-sdk/smithy-client";
+import {
+  INITIAL_RETRY_TOKENS,
+  RETRY_COST,
+  TIMEOUT_RETRY_COST,
+  NO_RETRY_INCREMENT
+} from "./constants";
+
+export const defaultRetryQuota = (): RetryQuota => {
+  const MAX_CAPACITY = INITIAL_RETRY_TOKENS;
+  let availableCapacity = INITIAL_RETRY_TOKENS;
+
+  const getCapacityAmount = (error: SdkError) =>
+    error.name === "TimeoutError" ? TIMEOUT_RETRY_COST : RETRY_COST;
+
+  const hasRetryTokens = (error: SdkError) =>
+    getCapacityAmount(error) <= availableCapacity;
+
+  const retrieveRetryToken = (error: SdkError) => {
+    if (!hasRetryTokens(error)) {
+      // retryStrategy should stop retrying, and return last error
+      throw new Error("No retry token available");
+    }
+    const capacityAmount = getCapacityAmount(error);
+    availableCapacity -= capacityAmount;
+    return capacityAmount;
+  };
+
+  const releaseRetryToken = (capacityReleaseAmount?: number) => {
+    availableCapacity += capacityReleaseAmount ?? NO_RETRY_INCREMENT;
+    availableCapacity = Math.min(availableCapacity, MAX_CAPACITY);
+  };
+
+  return Object.freeze({
+    hasRetryTokens,
+    retrieveRetryToken,
+    releaseRetryToken
+  });
+};
diff --git a/packages/middleware-retry/src/defaultStrategy.ts b/packages/middleware-retry/src/defaultStrategy.ts
index 7b91e9563589..e0d8c6f68ec2 100644
--- a/packages/middleware-retry/src/defaultStrategy.ts
+++ b/packages/middleware-retry/src/defaultStrategy.ts
@@ -66,7 +66,7 @@ export interface StandardRetryStrategyOptions {
 export class StandardRetryStrategy implements RetryStrategy {
   private retryDecider: RetryDecider;
   private delayDecider: DelayDecider;
-  private retryQuota?: RetryQuota;
+  // private retryQuota?: RetryQuota;
 
   constructor(
     public readonly maxAttempts: number,
@@ -74,7 +74,7 @@ export class StandardRetryStrategy implements RetryStrategy {
   ) {
     this.retryDecider = options?.retryDecider ?? defaultRetryDecider;
     this.delayDecider = options?.delayDecider ?? defaultDelayDecider;
-    this.retryQuota = options?.retryQuota;
+    // this.retryQuota = options?.retryQuota;
   }
 
   private shouldRetry(error: SdkError, attempts: number) {

From 592b08c26cf2d46a7ff4374d716283c551bce602 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Tue, 9 Jun 2020 23:12:38 +0000
Subject: [PATCH 5/7] feat: add defaultRetryQuota to defaultStrategy

---
 .../src/defaultRetryQuota.spec.ts             | 32 ++++++++++---------
 .../middleware-retry/src/defaultRetryQuota.ts | 10 +++---
 .../middleware-retry/src/defaultStrategy.ts   | 19 ++++++++---
 3 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/packages/middleware-retry/src/defaultRetryQuota.spec.ts b/packages/middleware-retry/src/defaultRetryQuota.spec.ts
index eb6e8df7d639..975cdae4218d 100644
--- a/packages/middleware-retry/src/defaultRetryQuota.spec.ts
+++ b/packages/middleware-retry/src/defaultRetryQuota.spec.ts
@@ -1,4 +1,4 @@
-import { defaultRetryQuota } from "./defaultRetryQuota";
+import { getDefaultRetryQuota } from "./defaultRetryQuota";
 import { SdkError } from "@aws-sdk/smithy-client";
 import {
   INITIAL_RETRY_TOKENS,
@@ -15,10 +15,10 @@ describe("defaultRetryQuota", () => {
     }) as SdkError;
 
   const getDrainedRetryQuota = (targetCapacity: number, error: SdkError) => {
-    const retryQuota = defaultRetryQuota();
+    const retryQuota = getDefaultRetryQuota();
     let availableCapacity = INITIAL_RETRY_TOKENS;
     while (availableCapacity >= targetCapacity) {
-      retryQuota.retrieveRetryToken(error);
+      retryQuota.retrieveRetryTokens(error);
       availableCapacity -= targetCapacity;
     }
     return retryQuota;
@@ -28,11 +28,13 @@ describe("defaultRetryQuota", () => {
     describe("returns true if capacity is available", () => {
       it("when it's TimeoutError", () => {
         const timeoutError = getMockTimeoutError();
-        expect(defaultRetryQuota().hasRetryTokens(timeoutError)).toBe(true);
+        expect(getDefaultRetryQuota().hasRetryTokens(timeoutError)).toBe(true);
       });
 
       it("when it's not TimeoutError", () => {
-        expect(defaultRetryQuota().hasRetryTokens(getMockError())).toBe(true);
+        expect(getDefaultRetryQuota().hasRetryTokens(getMockError())).toBe(
+          true
+        );
       });
     });
 
@@ -58,13 +60,13 @@ describe("defaultRetryQuota", () => {
     describe("returns retry tokens amount if available", () => {
       it("when it's TimeoutError", () => {
         const timeoutError = getMockTimeoutError();
-        expect(defaultRetryQuota().retrieveRetryToken(timeoutError)).toBe(
+        expect(getDefaultRetryQuota().retrieveRetryTokens(timeoutError)).toBe(
           TIMEOUT_RETRY_COST
         );
       });
 
       it("when it's not TimeoutError", () => {
-        expect(defaultRetryQuota().retrieveRetryToken(getMockError())).toBe(
+        expect(getDefaultRetryQuota().retrieveRetryTokens(getMockError())).toBe(
           RETRY_COST
         );
       });
@@ -78,7 +80,7 @@ describe("defaultRetryQuota", () => {
           timeoutError
         );
         expect(() => {
-          retryQuota.retrieveRetryToken(timeoutError);
+          retryQuota.retrieveRetryTokens(timeoutError);
         }).toThrowError(new Error("No retry token available"));
       });
 
@@ -86,7 +88,7 @@ describe("defaultRetryQuota", () => {
         const error = getMockError();
         const retryQuota = getDrainedRetryQuota(RETRY_COST, error);
         expect(() => {
-          retryQuota.retrieveRetryToken(error);
+          retryQuota.retrieveRetryTokens(error);
         }).toThrowError(new Error("No retry token available"));
       });
     });
@@ -101,9 +103,9 @@ describe("defaultRetryQuota", () => {
       expect(retryQuota.hasRetryTokens(error)).toBe(false);
 
       // Release RETRY_COST tokens.
-      retryQuota.releaseRetryToken(RETRY_COST);
+      retryQuota.releaseRetryTokens(RETRY_COST);
       expect(retryQuota.hasRetryTokens(error)).toBe(true);
-      expect(retryQuota.retrieveRetryToken(error)).toBe(RETRY_COST);
+      expect(retryQuota.retrieveRetryTokens(error)).toBe(RETRY_COST);
       expect(retryQuota.hasRetryTokens(error)).toBe(false);
     });
 
@@ -118,7 +120,7 @@ describe("defaultRetryQuota", () => {
         RETRY_COST - (INITIAL_RETRY_TOKENS % RETRY_COST);
       while (tokensReleased < tokensToBeReleased) {
         expect(retryQuota.hasRetryTokens(error)).toBe(false);
-        retryQuota.releaseRetryToken();
+        retryQuota.releaseRetryTokens();
         tokensReleased += NO_RETRY_INCREMENT;
       }
       expect(retryQuota.hasRetryTokens(error)).toBe(true);
@@ -126,11 +128,11 @@ describe("defaultRetryQuota", () => {
 
     it("ensures availableCapacity is maxed at INITIAL_RETRY_TOKENS", () => {
       const error = getMockError();
-      const retryQuota = defaultRetryQuota();
+      const retryQuota = getDefaultRetryQuota();
 
       // release 100 tokens.
       [...Array(100).keys()].forEach(key => {
-        retryQuota.releaseRetryToken();
+        retryQuota.releaseRetryTokens();
       });
 
       // availableCapacity is still maxed at INITIAL_RETRY_TOKENS
@@ -138,7 +140,7 @@ describe("defaultRetryQuota", () => {
       [...Array(Math.floor(INITIAL_RETRY_TOKENS / RETRY_COST)).keys()].forEach(
         key => {
           expect(retryQuota.hasRetryTokens(error)).toBe(true);
-          retryQuota.retrieveRetryToken(error);
+          retryQuota.retrieveRetryTokens(error);
         }
       );
       expect(retryQuota.hasRetryTokens(error)).toBe(false);
diff --git a/packages/middleware-retry/src/defaultRetryQuota.ts b/packages/middleware-retry/src/defaultRetryQuota.ts
index a2044b02ffdd..20250d304b34 100644
--- a/packages/middleware-retry/src/defaultRetryQuota.ts
+++ b/packages/middleware-retry/src/defaultRetryQuota.ts
@@ -7,7 +7,7 @@ import {
   NO_RETRY_INCREMENT
 } from "./constants";
 
-export const defaultRetryQuota = (): RetryQuota => {
+export const getDefaultRetryQuota = (): RetryQuota => {
   const MAX_CAPACITY = INITIAL_RETRY_TOKENS;
   let availableCapacity = INITIAL_RETRY_TOKENS;
 
@@ -17,7 +17,7 @@ export const defaultRetryQuota = (): RetryQuota => {
   const hasRetryTokens = (error: SdkError) =>
     getCapacityAmount(error) <= availableCapacity;
 
-  const retrieveRetryToken = (error: SdkError) => {
+  const retrieveRetryTokens = (error: SdkError) => {
     if (!hasRetryTokens(error)) {
       // retryStrategy should stop retrying, and return last error
       throw new Error("No retry token available");
@@ -27,14 +27,14 @@ export const defaultRetryQuota = (): RetryQuota => {
     return capacityAmount;
   };
 
-  const releaseRetryToken = (capacityReleaseAmount?: number) => {
+  const releaseRetryTokens = (capacityReleaseAmount?: number) => {
     availableCapacity += capacityReleaseAmount ?? NO_RETRY_INCREMENT;
     availableCapacity = Math.min(availableCapacity, MAX_CAPACITY);
   };
 
   return Object.freeze({
     hasRetryTokens,
-    retrieveRetryToken,
-    releaseRetryToken
+    retrieveRetryTokens,
+    releaseRetryTokens
   });
 };
diff --git a/packages/middleware-retry/src/defaultStrategy.ts b/packages/middleware-retry/src/defaultStrategy.ts
index e0d8c6f68ec2..8c7e467b784d 100644
--- a/packages/middleware-retry/src/defaultStrategy.ts
+++ b/packages/middleware-retry/src/defaultStrategy.ts
@@ -12,6 +12,7 @@ import {
   FinalizeHandlerArguments,
   RetryStrategy
 } from "@aws-sdk/types";
+import { getDefaultRetryQuota } from "./defaultRetryQuota";
 
 /**
  * Determines whether an error is retryable based on the number of retries
@@ -46,12 +47,12 @@ export interface RetryQuota {
    * returns token amount from the retry quota bucket.
    * throws error is retry tokens are not available.
    */
-  retrieveRetryToken: (error: SdkError) => number;
+  retrieveRetryTokens: (error: SdkError) => number;
 
   /**
    * releases tokens back to the retry quota.
    */
-  releaseRetryToken: (releaseCapacityAmount?: number) => void;
+  releaseRetryTokens: (releaseCapacityAmount?: number) => void;
 }
 
 /**
@@ -66,7 +67,7 @@ export interface StandardRetryStrategyOptions {
 export class StandardRetryStrategy implements RetryStrategy {
   private retryDecider: RetryDecider;
   private delayDecider: DelayDecider;
-  // private retryQuota?: RetryQuota;
+  private retryQuota: RetryQuota;
 
   constructor(
     public readonly maxAttempts: number,
@@ -74,22 +75,29 @@ export class StandardRetryStrategy implements RetryStrategy {
   ) {
     this.retryDecider = options?.retryDecider ?? defaultRetryDecider;
     this.delayDecider = options?.delayDecider ?? defaultDelayDecider;
-    // this.retryQuota = options?.retryQuota;
+    this.retryQuota = options?.retryQuota ?? getDefaultRetryQuota();
   }
 
   private shouldRetry(error: SdkError, attempts: number) {
-    return attempts < this.maxAttempts && this.retryDecider(error);
+    return (
+      attempts < this.maxAttempts &&
+      this.retryDecider(error) &&
+      this.retryQuota.hasRetryTokens(error)
+    );
   }
 
   async retry<Input extends object, Ouput extends MetadataBearer>(
     next: FinalizeHandler<Input, Ouput>,
     args: FinalizeHandlerArguments<Input>
   ) {
+    let retryTokenAmount;
     let attempts = 0;
     let totalDelay = 0;
     while (true) {
       try {
         const { response, output } = await next(args);
+
+        this.retryQuota.releaseRetryTokens(retryTokenAmount);
         output.$metadata.attempts = attempts + 1;
         output.$metadata.totalRetryDelay = totalDelay;
 
@@ -97,6 +105,7 @@ export class StandardRetryStrategy implements RetryStrategy {
       } catch (err) {
         attempts++;
         if (this.shouldRetry(err as SdkError, attempts)) {
+          retryTokenAmount = this.retryQuota.retrieveRetryTokens(err);
           const delay = this.delayDecider(
             isThrottlingError(err)
               ? THROTTLING_RETRY_DELAY_BASE

From e106e630d001edc7451eaa82815528af38874c2f Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Wed, 10 Jun 2020 01:04:52 +0000
Subject: [PATCH 6/7] test: defaultRetryQuota in defaultStrategy

---
 .../src/defaultStrategy.spec.ts               | 236 ++++++++++++++----
 1 file changed, 191 insertions(+), 45 deletions(-)

diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/defaultStrategy.spec.ts
index 0b9207a34eeb..c4ca6a00030a 100644
--- a/packages/middleware-retry/src/defaultStrategy.spec.ts
+++ b/packages/middleware-retry/src/defaultStrategy.spec.ts
@@ -6,6 +6,7 @@ import { isThrottlingError } from "@aws-sdk/service-error-classification";
 import { defaultDelayDecider } from "./delayDecider";
 import { defaultRetryDecider } from "./retryDecider";
 import { StandardRetryStrategy } from "./defaultStrategy";
+import { getDefaultRetryQuota } from "./defaultRetryQuota";
 
 jest.mock("@aws-sdk/service-error-classification", () => ({
   isThrottlingError: jest.fn().mockReturnValue(true)
@@ -19,8 +20,80 @@ jest.mock("./retryDecider", () => ({
   defaultRetryDecider: jest.fn().mockReturnValue(true)
 }));
 
+jest.mock("./defaultRetryQuota", () => {
+  const mockDefaultRetryQuota = {
+    hasRetryTokens: jest.fn().mockReturnValue(true),
+    retrieveRetryTokens: jest.fn().mockReturnValue(1),
+    releaseRetryTokens: jest.fn()
+  };
+  return { getDefaultRetryQuota: () => mockDefaultRetryQuota };
+});
+
 describe("defaultStrategy", () => {
-  const maxAttempts = 2;
+  const maxAttempts = 3;
+
+  const mockSuccessfulOperation = (maxAttempts: number, response?: string) => {
+    const next = jest.fn().mockResolvedValueOnce({
+      response,
+      output: { $metadata: {} }
+    });
+
+    const retryStrategy = new StandardRetryStrategy(maxAttempts);
+    return retryStrategy.retry(next, {} as any);
+  };
+
+  const mockFailedOperation = async (maxAttempts: number, error?: Error) => {
+    const mockError = error ?? new Error("mockError");
+    const next = jest.fn().mockRejectedValue(mockError);
+
+    const retryStrategy = new StandardRetryStrategy(maxAttempts);
+    try {
+      await retryStrategy.retry(next, {} as any);
+    } catch (error) {
+      expect(error).toStrictEqual(mockError);
+    }
+  };
+
+  const mockSuccessAfterOneFail = (
+    maxAttempts: number,
+    error?: Error,
+    response?: string
+  ) => {
+    const mockError = error ?? new Error("mockError");
+    const mockResponse = {
+      response,
+      output: { $metadata: {} }
+    };
+
+    const next = jest
+      .fn()
+      .mockRejectedValueOnce(mockError)
+      .mockResolvedValueOnce(mockResponse);
+
+    const retryStrategy = new StandardRetryStrategy(maxAttempts);
+    return retryStrategy.retry(next, {} as any);
+  };
+
+  const mockSuccessAfterTwoFails = (
+    maxAttempts: number,
+    error?: Error,
+    response?: string
+  ) => {
+    const mockError = error ?? new Error("mockError");
+    const mockResponse = {
+      response,
+      output: { $metadata: {} }
+    };
+
+    const next = jest
+      .fn()
+      .mockRejectedValueOnce(mockError)
+      .mockRejectedValueOnce(mockError)
+      .mockResolvedValueOnce(mockResponse);
+
+    const retryStrategy = new StandardRetryStrategy(maxAttempts);
+    return retryStrategy.retry(next, {} as any);
+  };
 
   afterEach(() => {
     jest.clearAllMocks();
@@ -80,19 +153,8 @@ describe("defaultStrategy", () => {
     ) => {
       (isThrottlingError as jest.Mock).mockReturnValueOnce(mockThrottlingError);
 
-      const mockError = new Error("mockError");
-      const mockResponse = {
-        response: "mockResponse",
-        output: { $metadata: {} }
-      };
-
-      const next = jest
-        .fn()
-        .mockRejectedValueOnce(mockError)
-        .mockResolvedValueOnce(mockResponse);
-
-      const retryStrategy = new StandardRetryStrategy(maxAttempts);
-      await retryStrategy.retry(next, {} as any);
+      const mockError = new Error();
+      await mockSuccessAfterOneFail(maxAttempts, mockError);
 
       expect(isThrottlingError as jest.Mock).toHaveBeenCalledTimes(1);
       expect(isThrottlingError as jest.Mock).toHaveBeenCalledWith(mockError);
@@ -112,19 +174,103 @@ describe("defaultStrategy", () => {
     });
   });
 
-  describe("should not retry", () => {
-    it("when the handler completes successfully", async () => {
-      const mockResponse = {
-        response: "mockResponse",
-        output: { $metadata: {} }
-      };
+  describe("retryQuota", () => {
+    describe("hasRetryTokens", () => {
+      it("not called on successful operation", async () => {
+        const { hasRetryTokens } = getDefaultRetryQuota();
+        await mockSuccessfulOperation(maxAttempts);
+        expect(hasRetryTokens).not.toHaveBeenCalled();
+      });
 
-      const next = jest.fn().mockResolvedValueOnce(mockResponse);
+      it("called once in case of single failure", async () => {
+        const { hasRetryTokens } = getDefaultRetryQuota();
+        await mockSuccessAfterOneFail(maxAttempts);
+        expect(hasRetryTokens).toHaveBeenCalledTimes(1);
+      });
 
-      const retryStrategy = new StandardRetryStrategy(maxAttempts);
-      const { response, output } = await retryStrategy.retry(next, {} as any);
+      it("called once on each retry request", async () => {
+        const { hasRetryTokens } = getDefaultRetryQuota();
+        await mockFailedOperation(maxAttempts);
+        expect(hasRetryTokens).toHaveBeenCalledTimes(maxAttempts - 1);
+      });
+    });
+
+    describe("releaseRetryTokens", () => {
+      it("called once without param on successful operation", async () => {
+        const { releaseRetryTokens } = getDefaultRetryQuota();
+        await mockSuccessfulOperation(maxAttempts);
+        expect(releaseRetryTokens).toHaveBeenCalledTimes(1);
+        expect(releaseRetryTokens).toHaveBeenCalledWith(undefined);
+      });
 
-      expect(response).toStrictEqual(mockResponse.response);
+      it("called once with retryTokenAmount in case of single failure", async () => {
+        const retryTokens = 15;
+        const {
+          releaseRetryTokens,
+          retrieveRetryTokens
+        } = getDefaultRetryQuota();
+        (retrieveRetryTokens as jest.Mock).mockReturnValueOnce(retryTokens);
+
+        await mockSuccessAfterOneFail(maxAttempts);
+        expect(releaseRetryTokens).toHaveBeenCalledTimes(1);
+        expect(releaseRetryTokens).toHaveBeenCalledWith(retryTokens);
+      });
+
+      it("called once with second retryTokenAmount in case of two failures", async () => {
+        const retryTokensFirst = 15;
+        const retryTokensSecond = 30;
+
+        const {
+          releaseRetryTokens,
+          retrieveRetryTokens
+        } = getDefaultRetryQuota();
+
+        (retrieveRetryTokens as jest.Mock)
+          .mockReturnValueOnce(retryTokensFirst)
+          .mockReturnValueOnce(retryTokensSecond);
+
+        await mockSuccessAfterTwoFails(maxAttempts);
+        expect(releaseRetryTokens).toHaveBeenCalledTimes(1);
+        expect(releaseRetryTokens).toHaveBeenCalledWith(retryTokensSecond);
+      });
+
+      it("not called on unsuccessful operation", async () => {
+        const { releaseRetryTokens } = getDefaultRetryQuota();
+        await mockFailedOperation(maxAttempts);
+        expect(releaseRetryTokens).not.toHaveBeenCalled();
+      });
+    });
+
+    describe("retrieveRetryTokens", () => {
+      it("not called on successful operation", async () => {
+        const { retrieveRetryTokens } = getDefaultRetryQuota();
+        await mockSuccessfulOperation(maxAttempts);
+        expect(retrieveRetryTokens).not.toHaveBeenCalled();
+      });
+
+      it("called once in case of single failure", async () => {
+        const { retrieveRetryTokens } = getDefaultRetryQuota();
+        await mockSuccessAfterOneFail(maxAttempts);
+        expect(retrieveRetryTokens).toHaveBeenCalledTimes(1);
+      });
+
+      it("called once on each retry request", async () => {
+        const { retrieveRetryTokens } = getDefaultRetryQuota();
+        await mockFailedOperation(maxAttempts);
+        expect(retrieveRetryTokens).toHaveBeenCalledTimes(maxAttempts - 1);
+      });
+    });
+  });
+
+  describe("should not retry", () => {
+    it("when the handler completes successfully", async () => {
+      const mockResponse = "mockResponse";
+      const { response, output } = await mockSuccessfulOperation(
+        maxAttempts,
+        mockResponse
+      );
+
+      expect(response).toStrictEqual(mockResponse);
       expect(output.$metadata.attempts).toBe(1);
       expect(output.$metadata.totalRetryDelay).toBe(0);
       expect(defaultRetryDecider as jest.Mock).not.toHaveBeenCalled();
@@ -133,35 +279,35 @@ describe("defaultStrategy", () => {
 
     it("when retryDecider returns false", async () => {
       (defaultRetryDecider as jest.Mock).mockReturnValueOnce(false);
-
-      const mockError = new Error("mockError");
-      const next = jest.fn().mockRejectedValueOnce(mockError);
-
-      const retryStrategy = new StandardRetryStrategy(maxAttempts);
-      try {
-        await retryStrategy.retry(next, {} as any);
-      } catch (error) {
-        expect(error).toStrictEqual(mockError);
-      }
-
+      const mockError = new Error();
+      await mockFailedOperation(maxAttempts, mockError);
       expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledTimes(1);
       expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledWith(mockError);
     });
 
-    it("when the the maximum number of attempts is reached", async () => {
-      const mockError = new Error("mockError");
-      const next = jest.fn().mockRejectedValue(mockError);
-
-      const retryStrategy = new StandardRetryStrategy(maxAttempts);
-      try {
-        await retryStrategy.retry(next, {} as any);
-      } catch (error) {
-        expect(error).toStrictEqual(mockError);
-      }
+    it("when the maximum number of attempts is reached", async () => {
+      await mockFailedOperation(maxAttempts);
       expect(defaultRetryDecider as jest.Mock).toHaveBeenCalledTimes(
         maxAttempts - 1
       );
     });
+
+    it("when retryQuota.hasRetryTokens returns false", async () => {
+      const {
+        hasRetryTokens,
+        retrieveRetryTokens,
+        releaseRetryTokens
+      } = getDefaultRetryQuota();
+      (hasRetryTokens as jest.Mock).mockReturnValueOnce(false);
+
+      const mockError = new Error();
+      await mockFailedOperation(maxAttempts, mockError);
+
+      expect(hasRetryTokens).toHaveBeenCalledTimes(1);
+      expect(hasRetryTokens).toHaveBeenCalledWith(mockError);
+      expect(retrieveRetryTokens).not.toHaveBeenCalled();
+      expect(releaseRetryTokens).not.toHaveBeenCalled();
+    });
   });
 
   it("should delay equal to the value returned by delayDecider", async () => {

From dc9db8d14032d34aea379931b2f1dd273d6af255 Mon Sep 17 00:00:00 2001
From: Trivikram Kamat <16024985+trivikr@users.noreply.github.com>
Date: Wed, 10 Jun 2020 02:11:34 +0000
Subject: [PATCH 7/7] test: add options.retryQuota tests

---
 .../src/defaultStrategy.spec.ts               | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/packages/middleware-retry/src/defaultStrategy.spec.ts b/packages/middleware-retry/src/defaultStrategy.spec.ts
index c4ca6a00030a..579ce90478da 100644
--- a/packages/middleware-retry/src/defaultStrategy.spec.ts
+++ b/packages/middleware-retry/src/defaultStrategy.spec.ts
@@ -5,7 +5,7 @@ import {
 import { isThrottlingError } from "@aws-sdk/service-error-classification";
 import { defaultDelayDecider } from "./delayDecider";
 import { defaultRetryDecider } from "./retryDecider";
-import { StandardRetryStrategy } from "./defaultStrategy";
+import { StandardRetryStrategy, RetryQuota } from "./defaultStrategy";
 import { getDefaultRetryQuota } from "./defaultRetryQuota";
 
 jest.mock("@aws-sdk/service-error-classification", () => ({
@@ -146,6 +146,26 @@ describe("defaultStrategy", () => {
     });
   });
 
+  describe("retryQuota", () => {
+    it("sets getDefaultRetryQuota if options is undefined", () => {
+      const retryStrategy = new StandardRetryStrategy(maxAttempts);
+      expect(retryStrategy["retryQuota"]).toBe(getDefaultRetryQuota());
+    });
+
+    it("sets getDefaultRetryQuota if options.delayDecider undefined", () => {
+      const retryStrategy = new StandardRetryStrategy(maxAttempts, {});
+      expect(retryStrategy["retryQuota"]).toBe(getDefaultRetryQuota());
+    });
+
+    it("sets options.retryQuota if defined", () => {
+      const retryQuota = {} as RetryQuota;
+      const retryStrategy = new StandardRetryStrategy(maxAttempts, {
+        retryQuota
+      });
+      expect(retryStrategy["retryQuota"]).toBe(retryQuota);
+    });
+  });
+
   describe("delayBase passed to delayDecider", () => {
     const testDelayBasePassed = async (
       delayBaseToTest: number,