From afee43712b72f2ef2dc7c515639a91c96545590f Mon Sep 17 00:00:00 2001
From: Chase Coalwell <chasecoa@amazon.com>
Date: Wed, 27 Nov 2019 14:09:13 -0800
Subject: [PATCH] Migrate acceptsHeaderMiddleware (#473)

---
 .../typescript/codegen/AddBuiltinPlugins.java |  5 ++
 .../src/index.spec.ts                         | 28 +++++------
 .../middleware-sdk-api-gateway/src/index.ts   | 46 +++++++++++++++----
 3 files changed, 52 insertions(+), 27 deletions(-)

diff --git a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddBuiltinPlugins.java b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddBuiltinPlugins.java
index e8738d88ba43..6838f825ce9b 100644
--- a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddBuiltinPlugins.java
+++ b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddBuiltinPlugins.java
@@ -52,6 +52,11 @@ public List<RuntimeClientPlugin> getClientPlugins() {
                 RuntimeClientPlugin.builder()
                         .withConventions(TypeScriptDependency.MIDDLEWARE_CONTENT_LENGTH.dependency, "ContentLength",
                                          HAS_MIDDLEWARE)
+                        .build(),
+                RuntimeClientPlugin.builder()
+                        .withConventions("@aws-sdk/middleware-sdk-api-gateway", "^0.1.0-preview.5",
+                                        "AcceptsHeader", HAS_MIDDLEWARE)
+                        .servicePredicate((m,s) -> s.getId().getName().equals("BackplaneControlService"))
                         .build()
         );
     }
diff --git a/packages/middleware-sdk-api-gateway/src/index.spec.ts b/packages/middleware-sdk-api-gateway/src/index.spec.ts
index 63c7c11b3258..a4661e1b2bce 100644
--- a/packages/middleware-sdk-api-gateway/src/index.spec.ts
+++ b/packages/middleware-sdk-api-gateway/src/index.spec.ts
@@ -1,31 +1,25 @@
-import { BuildHandler } from "@aws-sdk/types";
-import { acceptsHeader } from "./index";
+import { acceptsHeaderMiddleware } from "./index";
+import { HttpRequest } from "@aws-sdk/protocol-http";
 
-describe("acceptsHeader", () => {
-  const mockNextHandler = jest.fn();
-
-  const composedHandler: BuildHandler<any, any> = acceptsHeader(
-    mockNextHandler
-  );
+describe("acceptsHeaderMiddleware", () => {
+  const next = jest.fn();
 
   beforeEach(() => {
     jest.clearAllMocks();
   });
 
   it("sets Accepts header to application/json", async () => {
-    await composedHandler({
+    const handler = acceptsHeaderMiddleware()(next, {} as any);
+    await handler({
       input: {},
-      request: {
-        hostname: "foo.amazonaws.com",
-        path: "/",
-        protocol: "https:",
-        method: "GET",
+      request: new HttpRequest({
         headers: {}
-      }
+      })
     });
 
-    expect(mockNextHandler.mock.calls.length).toBe(1);
-    const { request } = mockNextHandler.mock.calls[0][0];
+    const { calls } = (next as any).mock;
+    expect(calls.length).toBe(1);
+    const { request } = next.mock.calls[0][0];
     expect(request.headers["accepts"]).toBe("application/json");
   });
 });
diff --git a/packages/middleware-sdk-api-gateway/src/index.ts b/packages/middleware-sdk-api-gateway/src/index.ts
index 9ba3115af9eb..1207c7e53a96 100644
--- a/packages/middleware-sdk-api-gateway/src/index.ts
+++ b/packages/middleware-sdk-api-gateway/src/index.ts
@@ -1,16 +1,42 @@
-import { BuildHandler, BuildHandlerArguments } from "@aws-sdk/types";
+import {
+  BuildHandler,
+  BuildHandlerArguments,
+  BuildHandlerOptions,
+  BuildHandlerOutput,
+  BuildMiddleware,
+  MetadataBearer,
+  Pluggable
+} from "@aws-sdk/types";
+import { HttpRequest } from "@aws-sdk/protocol-http";
 
-export function acceptsHeader(next: BuildHandler<any, any>) {
-  return async (args: BuildHandlerArguments<any, any>) => {
+export function acceptsHeaderMiddleware(): BuildMiddleware<any, any> {
+  return <Output extends MetadataBearer>(
+    next: BuildHandler<any, Output>
+  ): BuildHandler<any, Output> => async (
+    args: BuildHandlerArguments<any>
+  ): Promise<BuildHandlerOutput<Output>> => {
+    let request = { ...args.request };
+    if (HttpRequest.isInstance(request)) {
+      request.headers = {
+        ...request.headers,
+        accepts: "application/json"
+      };
+    }
     return next({
       ...args,
-      request: {
-        ...args.request,
-        headers: {
-          ...args.request.headers,
-          accepts: "application/json"
-        }
-      }
+      request
     });
   };
 }
+
+export const acceptsHeaderMiddlewareOptions: BuildHandlerOptions = {
+  step: "build",
+  tags: ["SET_ACCEPTS_HEADER", "ACCEPTS_HEADER"],
+  name: "acceptsHeaderMiddleware"
+};
+
+export const getAcceptsHeaderPlugin = (): Pluggable<any, any> => ({
+  applyToStack: clientStack => {
+    clientStack.add(acceptsHeaderMiddleware(), acceptsHeaderMiddlewareOptions);
+  }
+});