diff --git a/index.js b/index.js
index 6382895..19e27a5 100644
--- a/index.js
+++ b/index.js
@@ -22,9 +22,9 @@ module.exports.default = mock;
 Object.assign(mock, mm, {
   restore() {
     cluster.restore();
-    mm.restore();
+    mockAgent.restore(app);
     // return promise
-    return mockAgent.restore(app);
+    return mm.restore();
   },
 
   /**
diff --git a/lib/mock_agent.js b/lib/mock_agent.js
index 9eeba0d..4b544dd 100644
--- a/lib/mock_agent.js
+++ b/lib/mock_agent.js
@@ -8,12 +8,10 @@ module.exports = {
   getAgent(app) {
     if (!_global) {
       _global = getGlobalDispatcher();
-      if (typeof app?.httpclient?.getDispatcher === 'function') {
-        if (!app[APP_HTTPCLIENT_AGENT]) {
-          // save the raw dispatcher
-          app[APP_HTTPCLIENT_AGENT] = app.httpclient.getDispatcher();
-        }
-      }
+    }
+    if (!app?.httpclient?.[APP_HTTPCLIENT_AGENT] && typeof app?.httpclient?.getDispatcher === 'function') {
+      // save the raw dispatcher
+      app.httpclient[APP_HTTPCLIENT_AGENT] = app.httpclient.getDispatcher();
     }
     if (!_mockAgent) {
       _mockAgent = new MockAgent();
@@ -24,16 +22,14 @@ module.exports = {
     }
     return _mockAgent;
   },
-  async restore(app) {
+  restore(app) {
     if (!_mockAgent) return;
     if (_global) {
       setGlobalDispatcher(_global);
-      if (app?.[APP_HTTPCLIENT_AGENT]) {
-        app.httpclient.setDispatcher(app[APP_HTTPCLIENT_AGENT]);
-      }
     }
-    const agent = _mockAgent;
+    if (app?.httpclient?.[APP_HTTPCLIENT_AGENT]) {
+      app.httpclient.setDispatcher(app.httpclient[APP_HTTPCLIENT_AGENT]);
+    }
     _mockAgent = null;
-    await agent.close();
   },
 };