From b5024e600ff7c6e8b882ab7682ea8ca2d94843e0 Mon Sep 17 00:00:00 2001
From: fengmk2 <fengmk2@gmail.com>
Date: Mon, 4 Dec 2023 00:11:16 +0800
Subject: [PATCH] f

---
 src/HttpClient.ts  |   9 ++++
 test/index.test.ts | 104 +++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 109 insertions(+), 4 deletions(-)

diff --git a/src/HttpClient.ts b/src/HttpClient.ts
index 7c7300b1..58264f75 100644
--- a/src/HttpClient.ts
+++ b/src/HttpClient.ts
@@ -21,6 +21,7 @@ import {
   request as undiciRequest,
   Dispatcher,
   Agent,
+  getGlobalDispatcher,
 } from 'undici';
 import { FormData as FormDataNode } from 'formdata-node';
 import { FormDataEncoder } from 'form-data-encoder';
@@ -186,6 +187,14 @@ export class HttpClient extends EventEmitter {
     initDiagnosticsChannel();
   }
 
+  getDispatcher() {
+    return this.#dispatcher ?? getGlobalDispatcher();
+  }
+
+  setDispatcher(dispatcher: Dispatcher) {
+    this.#dispatcher = dispatcher;
+  }
+
   async request<T = any>(url: RequestURL, options?: RequestOptions) {
     return await this.#requestInternal<T>(url, options);
   }
diff --git a/test/index.test.ts b/test/index.test.ts
index 6fbe131b..630313f3 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -1,8 +1,8 @@
 import { strict as assert } from 'node:assert';
 import { parse as urlparse } from 'node:url';
 import { readFileSync } from 'node:fs';
-import { describe, it, beforeAll, afterAll } from 'vitest';
-import urllib from '../src';
+import { describe, it, beforeAll, afterAll, afterEach, beforeEach } from 'vitest';
+import urllib, { HttpClient } from '../src';
 import { MockAgent, setGlobalDispatcher, getGlobalDispatcher } from '../src';
 import { startServer } from './fixtures/server';
 import { readableToBytes } from './utils';
@@ -162,12 +162,12 @@ describe('index.test.ts', () => {
   describe('Mocking request', () => {
     let mockAgent: MockAgent;
     const globalAgent = getGlobalDispatcher();
-    beforeAll(() => {
+    beforeEach(() => {
       mockAgent = new MockAgent();
       setGlobalDispatcher(mockAgent);
     });
 
-    afterAll(async () => {
+    afterEach(async () => {
       setGlobalDispatcher(globalAgent);
       await mockAgent.close();
     });
@@ -275,5 +275,101 @@ describe('index.test.ts', () => {
 
       mockAgent.assertNoPendingInterceptors();
     });
+
+    it('should mocking intercept work on custom httpClient', async () => {
+      const httpClient = new HttpClient({
+        connect: {
+          timeout: 2000,
+        },
+      });
+      const oldAgent = httpClient.getDispatcher();
+      assert(oldAgent);
+      httpClient.setDispatcher(mockAgent);
+      const mockPool = mockAgent.get(_url.substring(0, _url.length - 1));
+      mockPool.intercept({
+        path: '/foo',
+        method: 'POST',
+      }).reply(400, {
+        message: 'mock 400 bad request',
+      });
+
+      mockPool.intercept({
+        path: '/bar',
+        method: 'GET',
+        query: {
+          q: '1',
+        },
+      }).reply(200, {
+        message: 'mock bar with q=1',
+      });
+
+      mockPool.intercept({
+        path: '/bar',
+        method: 'GET',
+        query: {
+          q: '2',
+        },
+      }).reply(200, {
+        message: 'mock bar with q=2',
+      });
+
+      mockPool.intercept({
+        path: /\.tgz$/,
+        method: 'GET',
+      }).reply(400, {
+        message: 'mock 400 bad request on tgz',
+      });
+
+      let response = await httpClient.request(`${_url}foo`, {
+        method: 'POST',
+        dataType: 'json',
+      });
+      assert.equal(response.status, 400);
+      assert.deepEqual(response.data, { message: 'mock 400 bad request' });
+
+      response = await httpClient.request(`${_url}bar?q=1`, {
+        method: 'GET',
+        dataType: 'json',
+      });
+      assert.equal(response.status, 200);
+      assert.deepEqual(response.data, { message: 'mock bar with q=1' });
+      response = await httpClient.request(`${_url}bar?q=2`, {
+        method: 'GET',
+        dataType: 'json',
+      });
+      assert.equal(response.status, 200);
+      assert.deepEqual(response.data, { message: 'mock bar with q=2' });
+
+      response = await httpClient.request(`${_url}download/foo.tgz`, {
+        method: 'GET',
+        dataType: 'json',
+      });
+      assert.equal(response.status, 400);
+      assert.deepEqual(response.data, { message: 'mock 400 bad request on tgz' });
+
+      // only intercept once
+      response = await httpClient.request(`${_url}download/bar.tgz`, {
+        method: 'GET',
+        dataType: 'json',
+      });
+      assert.equal(response.status, 200);
+      assert.equal(response.data.method, 'GET');
+
+      mockAgent.assertNoPendingInterceptors();
+
+      // should not work
+      httpClient.setDispatcher(oldAgent);
+      mockPool.intercept({
+        path: '/foo',
+        method: 'POST',
+      }).reply(400, {
+        message: 'mock 400 bad request',
+      });
+      response = await httpClient.request(`${_url}foo`, {
+        method: 'POST',
+        dataType: 'json',
+      });
+      assert.equal(response.status, 200);
+    });
   });
 });