Support multiple status codes in a @res decorator

lukeautry · Apr 14, 2021 · bdab059 · bdab059
1 parent 54dfca1
commit bdab059
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 22 deletions.
diff --git a/packages/cli/src/metadataGeneration/parameterGenerator.ts b/packages/cli/src/metadataGeneration/parameterGenerator.ts
@@ -31,7 +31,7 @@ export class ParameterGenerator {
       case 'Path':
         return [this.getPathParameter(this.parameter)];
       case 'Res':
-        return [this.getResParameter(this.parameter)];
+        return this.getResParameters(this.parameter);
       case 'Inject':
         return [];
       case 'UploadedFile':
@@ -56,7 +56,7 @@ export class ParameterGenerator {
     };
   }
 
-  private getResParameter(parameter: ts.ParameterDeclaration): Tsoa.ResParameter {
+  private getResParameters(parameter: ts.ParameterDeclaration): Tsoa.ResParameter[] {
     const parameterName = (parameter.name as ts.Identifier).text;
     const decorator = getNodeFirstDecoratorValue(this.parameter, this.current.typeChecker, ident => ident.text === 'Res') || parameterName;
     if (!decorator) {
@@ -74,33 +74,39 @@ export class ParameterGenerator {
     }
 
     const statusArgument = typeNode.typeArguments[0];
-    const statusArgumentType = this.current.typeChecker.getTypeAtLocation(statusArgument);
+    const bodyArgument = typeNode.typeArguments[1];
+
+    // support a union of status codes, all with the same response body
+    const statusArguments = (ts.isUnionTypeNode(statusArgument) ? [...statusArgument.types] : [statusArgument]);
+    const statusArgumentTypes = statusArguments.map(a => this.current.typeChecker.getTypeAtLocation(a));
 
     const isNumberLiteralType = (tsType: ts.Type): tsType is ts.NumberLiteralType => {
       // eslint-disable-next-line no-bitwise
       return (tsType.getFlags() & ts.TypeFlags.NumberLiteral) !== 0;
     };
 
-    if (!isNumberLiteralType(statusArgumentType)) {
-      throw new GenerateMetadataError('@Res() requires the type to be TsoaResponse<HTTPStatusCode, ResBody>', parameter);
-    }
-
-    const status = String(statusArgumentType.value);
-
-    const type = new TypeResolver(typeNode.typeArguments[1], this.current, typeNode).resolve();
+    return statusArgumentTypes.map(statusArgumentType => {
+      if (!isNumberLiteralType(statusArgumentType)) {
+        throw new GenerateMetadataError('@Res() requires the type to be TsoaResponse<HTTPStatusCode, ResBody>', parameter);
+      }
 
-    return {
-      description: this.getParameterDescription(parameter) || '',
-      in: 'res',
-      name: status,
-      parameterName,
-      examples: this.getParameterExample(parameter, parameterName),
-      required: true,
-      type,
-      schema: type,
-      validators: {},
-      headers: getHeaderType(typeNode.typeArguments, 2, this.current),
-    };
+      const status = String(statusArgumentType.value);
+
+      const type = new TypeResolver(bodyArgument, this.current, typeNode).resolve();
+
+      return {
+        description: this.getParameterDescription(parameter) || '',
+        in: 'res',
+        name: status,
+        parameterName,
+        examples: this.getParameterExample(parameter, parameterName),
+        required: true,
+        type,
+        schema: type,
+        validators: {},
+        headers: getHeaderType(typeNode.typeArguments, 2, this.current),
+      };
+    });
   }
 
   private getBodyPropParameter(parameter: ts.ParameterDeclaration): Tsoa.Parameter {

diff --git a/tests/fixtures/controllers/getController.ts b/tests/fixtures/controllers/getController.ts
@@ -238,6 +238,14 @@ export class GetTestController extends Controller {
       value: 'success',
     };
   }
+
+  /**
+   * @param res The alternate response
+   */
+   @Get('MultipleStatusCodeRes')
+   public async multipleStatusCodeRes(@Res() res: TsoaResponse<400 | 500, TestModel, { 'custom-header': string }>, @Query('statusCode') statusCode: 400 | 500): Promise<void> {
+     res?.(statusCode, new ModelService().getModel(), { 'custom-header': 'hello' });
+   }
 }
 
 export interface ErrorResponse {

diff --git a/tests/integration/express-server.spec.ts b/tests/integration/express-server.spec.ts
@@ -141,6 +141,20 @@ describe('Express Server', () => {
     );
   });
 
+  [400, 500].forEach(statusCode =>
+    it('Should support multiple status codes with the same @Res structure', () => {
+      return verifyGetRequest(
+        basePath + `/GetTest/MultipleStatusCodeRes?statusCode=${statusCode}`,
+        (err, res) => {
+          const model = res.body as TestModel;
+          expect(model.id).to.equal(1);
+          expect(res.get('custom-header')).to.eq('hello');
+        },
+        statusCode,
+      );
+    })
+  );
+
   it('Should not modify the response after headers sent', () => {
     return verifyGetRequest(
       basePath + '/GetTest/MultipleRes',

diff --git a/tests/integration/hapi-server.spec.ts b/tests/integration/hapi-server.spec.ts
@@ -1097,6 +1097,20 @@ describe('Hapi Server', () => {
       );
     });
 
+    [400, 500].forEach(statusCode =>
+      it('Should support multiple status codes with the same @Res structure', () => {
+        return verifyGetRequest(
+          basePath + `/GetTest/MultipleStatusCodeRes?statusCode=${statusCode}`,
+          (err, res) => {
+            const model = res.body as TestModel;
+            expect(model.id).to.equal(1);
+            expect(res.get('custom-header')).to.eq('hello');
+          },
+          statusCode,
+        );
+      })
+    );
+
     it('Should not modify the response after headers sent', () => {
       return verifyGetRequest(
         basePath + '/GetTest/MultipleRes',

diff --git a/tests/integration/koa-server.spec.ts b/tests/integration/koa-server.spec.ts
@@ -1086,6 +1086,20 @@ describe('Koa Server', () => {
       );
     });
 
+    [400, 500].forEach(statusCode =>
+      it('Should support multiple status codes with the same @Res structure', () => {
+        return verifyGetRequest(
+          basePath + `/GetTest/MultipleStatusCodeRes?statusCode=${statusCode}`,
+          (err, res) => {
+            const model = res.body as TestModel;
+            expect(model.id).to.equal(1);
+            expect(res.get('custom-header')).to.eq('hello');
+          },
+          statusCode,
+        );
+      })
+    );
+
     it('Should not modify the response after headers sent', () => {
       return verifyGetRequest(
         basePath + '/GetTest/MultipleRes',