fix: should reset agent on restore

eggjs · Jul 4, 2024 · 5c5da6f · 5c5da6f
1 parent e2a4b26
commit 5c5da6f
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 12 deletions.
diff --git a/app/extend/agent.js b/app/extend/agent.js
@@ -25,8 +25,8 @@ module.exports = {
     return mockAgent.getAgent(this);
   },
 
-  mockAgentRestore() {
-    return mockAgent.restore(this);
+  async mockAgentRestore() {
+    await mockAgent.restore();
   },
 
   /**

diff --git a/app/extend/application.js b/app/extend/application.js
@@ -306,8 +306,8 @@ module.exports = {
     return mockAgent.getAgent(this);
   },
 
-  mockAgentRestore() {
-    return mockAgent.restore(this);
+  async mockAgentRestore() {
+    await mockAgent.restore();
   },
 
   /**

diff --git a/index.js b/index.js
@@ -20,11 +20,10 @@ module.exports.default = mock;
 
 // inherit & extends mm
 Object.assign(mock, mm, {
-  restore() {
+  async restore() {
     cluster.restore();
-    mockAgent.restore(app);
-    // return promise
-    return mm.restore();
+    await mockAgent.restore();
+    return await mm.restore();
   },
 
   /**

diff --git a/lib/app.js b/lib/app.js
@@ -192,7 +192,7 @@ module.exports = function(options) {
       // don't delegate properties on MockApplication
       if (MOCK_APP_METHOD.includes(prop)) return getProperty(target, prop);
       if (!target[APP_INIT]) throw new Error(`can't get ${prop} before ready`);
-      // it's asyncrounus when agent and app are loading,
+      // it's asynchronous when agent and app are loading,
       // so should get the properties after loader ready
       debug('proxy handler.get %s', prop);
       return target._app[prop];

diff --git a/lib/mock_agent.js b/lib/mock_agent.js
@@ -1,8 +1,12 @@
+const { debuglog } = require('util');
 const { MockAgent, setGlobalDispatcher, getGlobalDispatcher } = require('urllib');
 
+const debug = debuglog('egg-mock:lib:mock_agent');
+
 let _mockAgent;
 let _global;
 const APP_HTTPCLIENT_AGENT = Symbol('app.httpclient.agent');
+const httpClients = [];
 
 module.exports = {
   getAgent(app) {
@@ -12,6 +16,8 @@ module.exports = {
     if (!app?.httpclient?.[APP_HTTPCLIENT_AGENT] && typeof app?.httpclient?.getDispatcher === 'function') {
       // save the raw dispatcher
       app.httpclient[APP_HTTPCLIENT_AGENT] = app.httpclient.getDispatcher();
+      httpClients.push(app.httpclient);
+      debug('add new httpClient, size: %d', httpClients.length);
     }
     if (!_mockAgent) {
       _mockAgent = new MockAgent();
@@ -22,14 +28,17 @@ module.exports = {
     }
     return _mockAgent;
   },
-  restore(app) {
+  async restore() {
     if (!_mockAgent) return;
     if (_global) {
       setGlobalDispatcher(_global);
     }
-    if (app?.httpclient?.[APP_HTTPCLIENT_AGENT]) {
-      app.httpclient.setDispatcher(app.httpclient[APP_HTTPCLIENT_AGENT]);
+    for (const httpClient of httpClients) {
+      httpClient.setDispatcher(httpClient[APP_HTTPCLIENT_AGENT]);
     }
+    debug('restore httpClient, size: %d', httpClients.length);
+    const agent = _mockAgent;
     _mockAgent = null;
+    await agent.close();
   },
 };
diff --git a/test/mock_httpclient_next_h2.test.js b/test/mock_httpclient_next_h2.test.js
@@ -192,6 +192,17 @@ describe('test/mock_httpclient_next_h2.test.js', () => {
       .expect(200);
   });
 
+  it('should auto restore after each case', async () => {
+    app.mockCsrf();
+    await request(server)
+      .get('/urllib')
+      .expect({
+        get: 'url get',
+        post: 'url post',
+      })
+      .expect(200);
+  });
+
   it('should use first mock data on duplicate url mock', async () => {
     app.mockCsrf();
     app.mockHttpclient(url, 'post', {