feat: Allow overriding status codes in response (#12)

typeofweb · web-flow · commit 23741d2182e0 · 2021-06-21T17:21:04.000+02:00
diff --git a/__tests__/router.test.ts b/__tests__/router.test.ts
@@ -1,6 +1,6 @@
 import Crypto from 'crypto';
 
-import { createApp } from '../src';
+import { createApp, HttpStatusCode } from '../src';
 import { unseal } from '../src/utils/encryptCookies';
 
 declare module '../src' {
@@ -10,107 +10,129 @@ declare module '../src' {
 }
 
 describe('router', () => {
-  it('should set cookies', async () => {
-    const app = createApp({ cookies: { secret: 'a'.repeat(32) } }).route({
-      path: '/users',
-      method: 'get',
-      validation: {},
-      handler: async (_request, t) => {
-        await t.setCookie('test', 'testowa', { encrypted: false });
-        return null;
-      },
-    });
+  describe('cookies', () => {
+    it('should set cookies', async () => {
+      const app = createApp({ cookies: { secret: 'a'.repeat(32) } }).route({
+        path: '/users',
+        method: 'get',
+        validation: {},
+        handler: async (_request, t) => {
+          await t.setCookie('test', 'testowa', { encrypted: false });
+          return null;
+        },
+      });
 
-    const result = await app.inject({
-      method: 'get',
-      path: '/users',
+      const result = await app.inject({
+        method: 'get',
+        path: '/users',
+      });
+
+      expect(result.headers['set-cookie']).toEqual([`test=testowa; Path=/; HttpOnly; Secure; SameSite=Lax`]);
     });
 
-    expect(result.headers['set-cookie']).toEqual([`test=testowa; Path=/; HttpOnly; Secure; SameSite=Lax`]);
-  });
+    it('should set encrypted cookies', async () => {
+      const secret = Crypto.randomBytes(22).toString('base64');
+      const app = createApp({ cookies: { secret } }).route({
+        path: '/users',
+        method: 'get',
+        validation: {},
+        handler: async (_request, t) => {
+          await t.setCookie('test', 'testowa', { encrypted: true });
+          return null;
+        },
+      });
 
-  it('should set encrypted cookies', async () => {
-    const secret = Crypto.randomBytes(22).toString('base64');
-    const app = createApp({ cookies: { secret } }).route({
-      path: '/users',
-      method: 'get',
-      validation: {},
-      handler: async (_request, t) => {
-        await t.setCookie('test', 'testowa', { encrypted: true });
-        return null;
-      },
-    });
+      const result = await app.inject({
+        method: 'get',
+        path: '/users',
+      });
 
-    const result = await app.inject({
-      method: 'get',
-      path: '/users',
+      expect(result.headers['set-cookie']).toHaveLength(1);
+      const [key, val] = (result.headers['set-cookie'] as readonly [string])[0].split(';')[0]!.split('=');
+      expect(key).toEqual('test');
+      expect(await unseal({ sealed: val!, secret })).toEqual('testowa');
     });
 
-    expect(result.headers['set-cookie']).toHaveLength(1);
-    const [key, val] = (result.headers['set-cookie'] as readonly [string])[0].split(';')[0]!.split('=');
-    expect(key).toEqual('test');
-    expect(await unseal({ sealed: val!, secret })).toEqual('testowa');
-  });
+    it('should read all cookies', async () => {
+      const app = createApp({ cookies: { secret: 'a'.repeat(32) } }).route({
+        path: '/users',
+        method: 'get',
+        validation: {},
+        handler: (request) => {
+          expect(request.cookies['test']).toEqual('testowa');
+          expect(request.cookies['secret']).toEqual('testowa');
+          return null;
+        },
+      });
 
-  it('should read all cookies', async () => {
-    const app = createApp({ cookies: { secret: 'a'.repeat(32) } }).route({
-      path: '/users',
-      method: 'get',
-      validation: {},
-      handler: (request) => {
-        expect(request.cookies['test']).toEqual('testowa');
-        expect(request.cookies['secret']).toEqual('testowa');
-        return null;
-      },
+      await app.inject({
+        method: 'get',
+        path: '/users',
+        cookies: [
+          'test=testowa',
+          'secret=Fe26.2**n0Jf_bWyQEP8IGbXE2USt82PE9W_dGIvOSlLTeKvwnA*a_fJJWNvbhjvvy0yBg2APw*PzP5xlZ63c6EImBsgDaZ7A**LF1tHWLrtR1pckVC9moT-V2b8LVmE_NYWfsgYqYhZwk*i40hiIKMSNNpwQUmd2HalR9KEvODTfDRPLah2H_q2JqdPg3ecHW6PvnJTC2YAHGUiwvIERWqE5LvaSjCyULvbQ',
+        ],
+      });
+
+      expect.assertions(2);
     });
 
-    await app.inject({
-      method: 'get',
-      path: '/users',
-      cookies: [
-        'test=testowa',
-        'secret=Fe26.2**n0Jf_bWyQEP8IGbXE2USt82PE9W_dGIvOSlLTeKvwnA*a_fJJWNvbhjvvy0yBg2APw*PzP5xlZ63c6EImBsgDaZ7A**LF1tHWLrtR1pckVC9moT-V2b8LVmE_NYWfsgYqYhZwk*i40hiIKMSNNpwQUmd2HalR9KEvODTfDRPLah2H_q2JqdPg3ecHW6PvnJTC2YAHGUiwvIERWqE5LvaSjCyULvbQ',
-      ],
+    it('should not accept multiple params in path segment', () => {
+      const app = createApp({});
+
+      return expect(() =>
+        app.route({
+          path: '/currencies/:from-:to',
+          method: 'get',
+          validation: {
+            params: {} as any,
+          },
+
+          handler: () => {
+            return { message: 'OKEJ' };
+          },
+        }),
+      ).toThrowErrorMatchingInlineSnapshot(`"RouteValidationError: Each path segment can contain at most one param."`);
     });
 
-    expect.assertions(2);
+    it('should not accept regexes', () => {
+      const app = createApp({});
+
+      return expect(() =>
+        app.route({
+          path: '/currencies/:from(\\d+)',
+          method: 'get',
+          validation: {
+            params: {} as any,
+          },
+
+          handler: () => {
+            return { message: 'OKEJ' };
+          },
+        }),
+      ).toThrowErrorMatchingInlineSnapshot(
+        `"RouteValidationError: Don't use regular expressions in routes. Use validators instead."`,
+      );
+    });
   });
 
-  it('should not accept multiple params in path segment', () => {
-    const app = createApp({});
-
-    return expect(() =>
-      app.route({
-        path: '/currencies/:from-:to',
+  describe('status code', () => {
+    it('should allow overriding status codes', async () => {
+      const app = createApp({}).route({
+        path: '/users',
         method: 'get',
-        validation: {
-          params: {} as any,
-        },
-
-        handler: () => {
-          return { message: 'OKEJ' };
+        validation: {},
+        handler: async (req, t) => {
+          await t.setStatus(HttpStatusCode.Accepted);
+          return null;
         },
-      }),
-    ).toThrowErrorMatchingInlineSnapshot(`"RouteValidationError: Each path segment can contain at most one param."`);
-  });
-
-  it('should not accept regexes', () => {
-    const app = createApp({});
+      });
 
-    return expect(() =>
-      app.route({
-        path: '/currencies/:from(\\d+)',
+      const result = await app.inject({
         method: 'get',
-        validation: {
-          params: {} as any,
-        },
-
-        handler: () => {
-          return { message: 'OKEJ' };
-        },
-      }),
-    ).toThrowErrorMatchingInlineSnapshot(
-      `"RouteValidationError: Don't use regular expressions in routes. Use validators instead."`,
-    );
+        path: '/users',
+      });
+      expect(result.statusCode).toEqual(202);
+    });
   });
 });
diff --git a/src/modules/router.ts b/src/modules/router.ts
@@ -93,7 +93,9 @@ export const errorMiddleware =
   };
 
 type AsyncHandler = (
-  ...args: Parameters<Express.Handler>
+  req: Express.Request,
+  res: Express.Response<any, ExpressResponseLocals>,
+  next: Express.NextFunction,
 ) => Promise<ReturnType<Express.Handler>> | ReturnType<Express.Handler>;
 
 const finalErrorGuard = (h: AsyncHandler): AsyncHandler => {
@@ -241,28 +243,40 @@ export const routeToExpressHandler = <
     // eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- response is always Json or null
     server.events.emit(':response', result.value as Json | null);
 
+    const fallbackStatusCode =
+      result.value === null || result.value === undefined ? HttpStatusCode.NoContent : HttpStatusCode.OK;
+    const statusCode = res.locals[CUSTOM_STATUS_CODE] ?? fallbackStatusCode;
+
     if (result.value === null) {
-      res.status(204).end();
+      res.status(statusCode).end();
       return;
     } else if (result.value === undefined) {
       console.warn(
         'Handler returned `undefined` which usually means you forgot to `await` something. If you want an empty response, return `null` instead.',
       );
-      res.status(204).end();
+      res.status(statusCode).end();
       return;
     } else {
-      res.status(200).json(result.value);
+      res.status(statusCode).json(result.value);
       return;
     }
   };
 };
 
+export const CUSTOM_STATUS_CODE = Symbol('CUSTOM_STATUS_CODE');
+export interface ExpressResponseLocals {
+  /* eslint-disable functional/prefer-readonly-type -- these are writable */
+  [key: string]: any;
+  [CUSTOM_STATUS_CODE]?: HttpStatusCode;
+  /* eslint-enable functional/prefer-readonly-type */
+}
+
 function createRequestToolkitFor({
   appOptions,
   res,
 }: {
   readonly req: Express.Request;
-  readonly res: Express.Response;
+  readonly res: Express.Response<any, ExpressResponseLocals>;
   readonly appOptions: AppOptions;
 }): TypeOfWebRequestToolkit {
   const toolkit: TypeOfWebRequestToolkit = {
@@ -280,6 +294,9 @@ function createRequestToolkitFor({
       const cookieOptions = deepMerge(options, appOptions.cookies);
       res.clearCookie(name, cookieOptions);
     },
+    setStatus(statusCode: HttpStatusCode) {
+      res.locals[CUSTOM_STATUS_CODE] = statusCode;
+    },
   };
 
   return toolkit;
diff --git a/src/modules/shared.ts b/src/modules/shared.ts
@@ -1,7 +1,7 @@
 import type { TypeOfWebRequestMeta, TypeOfWebServerMeta, TypeOfWebEvents } from '..';
 import type { Callback, Json, MaybeAsync } from '../utils/types';
 import type { RequestId, ServerId } from '../utils/uniqueId';
-import type { HttpMethod } from './httpStatusCodes';
+import type { HttpMethod, HttpStatusCode } from './httpStatusCodes';
 import type { TypeOfWebPlugin } from './plugins';
 import type { ParseRouteParams } from './router';
 import type { SchemaRecord, TypeOfRecord } from './validation';
@@ -40,6 +40,7 @@ interface AppOptionsCookies extends SetCookieOptions {
 export interface TypeOfWebRequestToolkit {
   setCookie(name: string, value: string, options?: SetCookieOptions): MaybeAsync<void>;
   removeCookie(name: string, options?: SetCookieOptions): MaybeAsync<void>;
+  setStatus(statusCode: HttpStatusCode): MaybeAsync<void>;
 }
 
 export interface SetCookieOptions {
