From 4647b9761943cab24a6172735c7afb4b1d25745e Mon Sep 17 00:00:00 2001
From: Swain Molster <swain.molster@lifeomic.com>
Date: Tue, 23 Jul 2024 21:36:57 -0400
Subject: [PATCH] feat: support mocking response headers

---
 src/test-utils.test.ts | 38 ++++++++++++++++++++++++++++++++++++++
 src/test-utils.ts      | 38 ++++++++++++++++++++++++++------------
 2 files changed, 64 insertions(+), 12 deletions(-)

diff --git a/src/test-utils.test.ts b/src/test-utils.test.ts
index 26e0699..c424ce3 100644
--- a/src/test-utils.test.ts
+++ b/src/test-utils.test.ts
@@ -82,6 +82,44 @@ describe('createAPIMockingUtility', () => {
       expect(result.status).toStrictEqual(200);
       expect(result.data).toStrictEqual({ message: payload.message });
     });
+
+    test('static response with headers', async () => {
+      network.mock('GET /items/:id', {
+        status: 200,
+        data: { message: 'some-message' },
+        headers: {
+          'test-response-header': 'some-value',
+        },
+      });
+
+      const payload = { id: v4(), filter: v4() };
+      const result = await client.request('GET /items/:id', payload);
+
+      expect(result.status).toStrictEqual(200);
+      expect(result.data).toStrictEqual({ message: 'some-message' });
+      expect(result.headers).toMatchObject({
+        'test-response-header': 'some-value',
+      });
+    });
+
+    test('function response with headers', async () => {
+      network.mock('GET /items/:id', () => ({
+        status: 200,
+        data: { message: 'some-message' },
+        headers: {
+          'test-response-header': 'some-value',
+        },
+      }));
+
+      const payload = { id: v4(), filter: v4() };
+      const result = await client.request('GET /items/:id', payload);
+
+      expect(result.status).toStrictEqual(200);
+      expect(result.data).toStrictEqual({ message: 'some-message' });
+      expect(result.headers).toMatchObject({
+        'test-response-header': 'some-value',
+      });
+    });
   });
 
   test('mockOrdered', async () => {
diff --git a/src/test-utils.ts b/src/test-utils.ts
index 5ebd453..051a80e 100644
--- a/src/test-utils.ts
+++ b/src/test-utils.ts
@@ -1,4 +1,4 @@
-import { RequestHandler, rest } from 'msw';
+import { RequestHandler, ResponseTransformer, rest } from 'msw';
 /**
  * Important: we need to avoid import actual values from these imports:
  * - 'msw/node'
@@ -20,8 +20,16 @@ type MSWUsable = {
 };
 
 export type APIMockerResponse<T> =
-  | { status: 200; data: T }
-  | { status: 400 | 401 | 403 | 404 | 500; data: any };
+  | {
+      status: 200;
+      data: T;
+      headers?: Record<string, string>;
+    }
+  | {
+      status: 400 | 401 | 403 | 404 | 500;
+      data: any;
+      headers?: Record<string, string>;
+    };
 
 type MockRequestHandler<
   Endpoints extends RoughEndpoints,
@@ -112,13 +120,22 @@ export const createAPIMocker = <Endpoints extends RoughEndpoints>(
 
       server.use(
         rest[lowercaseMethod](`${baseUrl}${url}`, async (req, res, ctx) => {
-          const resolve = options.once ? res.once : res;
+          const toMSWResponse = (response: APIMockerResponse<any>) => {
+            const resolve = options.once ? res.once : res;
+            const transformers: ResponseTransformer[] = [
+              ctx.status(response.status),
+              ctx.json(response.data),
+            ];
+            if (response.headers) {
+              for (const [key, value] of Object.entries(response.headers)) {
+                transformers.push(ctx.set(key, value));
+              }
+            }
+            return resolve(...transformers);
+          };
 
           if (typeof handlerOrResponse !== 'function') {
-            return resolve(
-              ctx.status(handlerOrResponse.status),
-              ctx.json(handlerOrResponse.data),
-            );
+            return toMSWResponse(handlerOrResponse);
           }
 
           const mockRequest = {
@@ -151,10 +168,7 @@ export const createAPIMocker = <Endpoints extends RoughEndpoints>(
             });
           }
 
-          return resolve(
-            ctx.status(mockedResponse.status),
-            ctx.json(mockedResponse.data),
-          );
+          return toMSWResponse(mockedResponse);
         }),
       );