From 829faa04a43f07f6bf040986ebe5a7cf2a7e254d Mon Sep 17 00:00:00 2001
From: jannyHou <juehou@ca.ibm.com>
Date: Mon, 16 Apr 2018 13:54:22 -0400
Subject: [PATCH] poc: customize statusCode using response spec

---
 .../src/openapi-spec-builder.ts                    | 12 ++++++++++++
 packages/rest/src/internal-types.ts                |  6 +++++-
 packages/rest/src/sequence.ts                      |  2 +-
 packages/rest/src/writer.ts                        | 10 ++++++++++
 .../test/integration/http-handler.integration.ts   | 14 ++++++++++++++
 5 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/packages/openapi-spec-builder/src/openapi-spec-builder.ts b/packages/openapi-spec-builder/src/openapi-spec-builder.ts
index cce3a7e40c29..c67ee320ac0a 100644
--- a/packages/openapi-spec-builder/src/openapi-spec-builder.ts
+++ b/packages/openapi-spec-builder/src/openapi-spec-builder.ts
@@ -112,6 +112,18 @@ export class OpenApiSpecBuilder extends BuilderBase<OpenApiSpec> {
 
     return this.withOperation(verb, path, spec);
   }
+
+  withOperationReturningStringWithStatusCode(
+    verb: string,
+    path: string,
+    statusCode: number,
+    operationName?: string,
+  ): this {
+    const spec = anOperationSpec().withStringResponse(statusCode);
+    if (operationName) spec.withOperationName(operationName);
+
+    return this.withOperation(verb, path, spec);
+  }
 }
 
 /**
diff --git a/packages/rest/src/internal-types.ts b/packages/rest/src/internal-types.ts
index 74e9694e9c36..ee42732c7182 100644
--- a/packages/rest/src/internal-types.ts
+++ b/packages/rest/src/internal-types.ts
@@ -52,7 +52,11 @@ export type InvokeMethod = (
  * @param response The response the response to send to.
  * @param result The operation result to send.
  */
-export type Send = (response: ServerResponse, result: OperationRetval) => void;
+export type Send = (
+  response: ServerResponse,
+  result: OperationRetval,
+  route?: ResolvedRoute,
+) => void;
 
 /**
  * Reject the request with an error.
diff --git a/packages/rest/src/sequence.ts b/packages/rest/src/sequence.ts
index e57da1f97430..c96be8e707bb 100644
--- a/packages/rest/src/sequence.ts
+++ b/packages/rest/src/sequence.ts
@@ -108,7 +108,7 @@ export class DefaultSequence implements SequenceHandler {
       const result = await this.invoke(route, args);
 
       debug('%s result -', route.describe(), result);
-      this.send(res, result);
+      this.send(res, result, route);
     } catch (err) {
       this.reject(res, req, err);
     }
diff --git a/packages/rest/src/writer.ts b/packages/rest/src/writer.ts
index 9c012265cad4..1bbc8ed3bd6c 100644
--- a/packages/rest/src/writer.ts
+++ b/packages/rest/src/writer.ts
@@ -6,6 +6,7 @@
 import {ServerResponse as Response} from 'http';
 import {OperationRetval} from './internal-types';
 import {HttpError} from 'http-errors';
+import {ResolvedRoute} from './router/routing-table';
 import {Readable} from 'stream';
 
 /**
@@ -20,6 +21,7 @@ export function writeResultToResponse(
   response: Response,
   // result returned back from invoking controller method
   result: OperationRetval,
+  route?: ResolvedRoute,
 ): void {
   if (result) {
     if (result instanceof Readable || typeof result.pipe === 'function') {
@@ -48,6 +50,14 @@ export function writeResultToResponse(
         result = result.toString();
         break;
     }
+    if (route) {
+      const spec = route.spec;
+      const responsesCode = Object.keys(spec.responses || {});
+      if (responsesCode.length >= 1) {
+        const successStatusCode = responsesCode[0];
+        response.statusCode = parseInt(successStatusCode);
+      }
+    }
     response.write(result);
   }
   response.end();
diff --git a/packages/rest/test/integration/http-handler.integration.ts b/packages/rest/test/integration/http-handler.integration.ts
index c03936de4b4c..0782de108702 100644
--- a/packages/rest/test/integration/http-handler.integration.ts
+++ b/packages/rest/test/integration/http-handler.integration.ts
@@ -58,6 +58,12 @@ describe('HttpHandler', () => {
         .withOperationReturningString('get', '/hello', 'hello')
         .withOperationReturningString('get', '/bye', 'bye')
         .withOperationReturningString('post', '/hello', 'postHello')
+        .withOperationReturningStringWithStatusCode(
+          'post',
+          '/createHello',
+          201,
+          'createHello',
+        )
         .build();
 
       class HelloController {
@@ -72,11 +78,19 @@ describe('HttpHandler', () => {
         public async postHello(): Promise<string> {
           return 'hello posted';
         }
+
+        public async createHello(): Promise<string> {
+          return 'hello created';
+        }
       }
 
       givenControllerClass(HelloController, spec);
     });
 
+    it('returns 201 for "POST /createHello"', () => {
+      return client.post('/createHello').expect(201, 'hello created');
+    });
+
     it('executes hello() for "GET /hello"', () => {
       return client.get('/hello').expect('hello');
     });