From 33b253631e59e06b95be9f122f817e86f7c00d5e Mon Sep 17 00:00:00 2001
From: Filip Skokan <panva.ip@gmail.com>
Date: Sun, 13 Mar 2022 19:08:40 +0100
Subject: [PATCH] feat(mock): allow matching the complete header list (#1275)

---
 lib/mock/mock-utils.js                | 11 ++++++
 test/mock-agent.js                    | 56 +++++++++++++++++++++++++++
 test/types/mock-interceptor.test-d.ts |  7 ++++
 types/mock-interceptor.d.ts           |  2 +-
 4 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/lib/mock/mock-utils.js b/lib/mock/mock-utils.js
index f9785f456be..0d4b86bba52 100644
--- a/lib/mock/mock-utils.js
+++ b/lib/mock/mock-utils.js
@@ -24,7 +24,18 @@ function matchValue (match, value) {
   return false
 }
 
+function lowerCaseEntries (headers) {
+  return Object.fromEntries(
+    Object.entries(headers).map(([headerName, headerValue]) => {
+      return [headerName.toLocaleLowerCase(), headerValue]
+    })
+  )
+}
+
 function matchHeaders (mockDispatch, headers) {
+  if (typeof mockDispatch.headers === 'function') {
+    return mockDispatch.headers(headers ? lowerCaseEntries(headers) : {})
+  }
   if (typeof mockDispatch.headers === 'undefined') {
     return true
   }
diff --git a/test/mock-agent.js b/test/mock-agent.js
index c04daba5508..db0beb52ebd 100644
--- a/test/mock-agent.js
+++ b/test/mock-agent.js
@@ -2293,3 +2293,59 @@ test('MockAgent - disableNetConnect should throw if dispatch not found by net co
     method: 'GET'
   }), new MockNotMatchedError(`Mock dispatch not matched for path '/foo': subsequent request to origin ${baseUrl} was not allowed (net.connect disabled)`))
 })
+
+test('MockAgent - headers function interceptor', async (t) => {
+  t.plan(7)
+
+  const server = createServer((req, res) => {
+    t.fail('should not be called')
+    t.end()
+    res.end('should not be called')
+  })
+  t.teardown(server.close.bind(server))
+
+  await promisify(server.listen.bind(server))(0)
+
+  const baseUrl = `http://localhost:${server.address().port}`
+
+  const mockAgent = new MockAgent()
+  setGlobalDispatcher(mockAgent)
+  t.teardown(mockAgent.close.bind(mockAgent))
+  const mockPool = mockAgent.get(baseUrl)
+
+  // Disable net connect so we can make sure it matches properly
+  mockAgent.disableNetConnect()
+
+  mockPool.intercept({
+    path: '/foo',
+    method: 'GET',
+    headers (headers) {
+      t.equal(typeof headers, 'object')
+      return !Object.keys(headers).includes('authorization')
+    }
+  }).reply(200, 'foo').times(2)
+
+  await t.rejects(request(`${baseUrl}/foo`, {
+    method: 'GET',
+    headers: {
+      Authorization: 'Bearer foo'
+    }
+  }), new MockNotMatchedError(`Mock dispatch not matched for headers '{"Authorization":"Bearer foo"}': subsequent request to origin ${baseUrl} was not allowed (net.connect disabled)`))
+
+  {
+    const { statusCode } = await request(`${baseUrl}/foo`, {
+      method: 'GET',
+      headers: {
+        foo: 'bar'
+      }
+    })
+    t.equal(statusCode, 200)
+  }
+
+  {
+    const { statusCode } = await request(`${baseUrl}/foo`, {
+      method: 'GET'
+    })
+    t.equal(statusCode, 200)
+  }
+})
diff --git a/test/types/mock-interceptor.test-d.ts b/test/types/mock-interceptor.test-d.ts
index a1c487ff211..1d04b8385ab 100644
--- a/test/types/mock-interceptor.test-d.ts
+++ b/test/types/mock-interceptor.test-d.ts
@@ -61,3 +61,10 @@ import { MockInterceptor, MockScope } from '../../types/mock-interceptor'
   // times
   expectAssignable<MockScope>(mockScope.times(2))
 }
+
+{
+  const mockPool: MockPool = new MockAgent().get('')
+  mockPool.intercept({ path: '', method: 'GET', headers: () => true })
+  mockPool.intercept({ path: '', method: 'GET', headers: () => false })
+  mockPool.intercept({ path: '', method: 'GET', headers: (headers) => Object.keys(headers).includes('authorization') })
+}
diff --git a/types/mock-interceptor.d.ts b/types/mock-interceptor.d.ts
index eb7f531bcf2..0166b1f1db3 100644
--- a/types/mock-interceptor.d.ts
+++ b/types/mock-interceptor.d.ts
@@ -49,7 +49,7 @@ declare namespace MockInterceptor {
     /** Body to intercept on. */
     body?: string | RegExp | ((body: string) => boolean);
     /** Headers to intercept on. */
-    headers?: Record<string, string | RegExp | ((body: string) => boolean)>;
+    headers?: Record<string, string | RegExp | ((body: string) => boolean)> | ((headers: Record<string, string>) => boolean);
   }
   export interface MockDispatch<TData extends object = object, TError extends Error = Error> extends Options {
     times: number | null;