1 files changed, 59 insertions, 0 deletions
diff --git a/lib/mock/mock-pool.js b/lib/mock/mock-pool.js
new file mode 100644
index 0000000..0a3a7cd
--- /dev/null
+++ b/lib/mock/mock-pool.js
@@ -0,0 +1,59 @@
+'use strict'
+
+const { promisify } = require('util')
+const Pool = require('../pool')
+const { buildMockDispatch } = require('./mock-utils')
+const {
+  kDispatches,
+  kMockAgent,
+  kClose,
+  kOriginalClose,
+  kOrigin,
+  kOriginalDispatch,
+  kConnected
+} = require('./mock-symbols')
+const { MockInterceptor } = require('./mock-interceptor')
+const Symbols = require('../core/symbols')
+const { InvalidArgumentError } = require('../core/errors')
+
+/**
+ * MockPool provides an API that extends the Pool to influence the mockDispatches.
+ */
+class MockPool extends Pool {
+  constructor (origin, opts) {
+    super(origin, opts)
+
+    if (!opts || !opts.agent || typeof opts.agent.dispatch !== 'function') {
+      throw new InvalidArgumentError('Argument opts.agent must implement Agent')
+    }
+
+    this[kMockAgent] = opts.agent
+    this[kOrigin] = origin
+    this[kDispatches] = []
+    this[kConnected] = 1
+    this[kOriginalDispatch] = this.dispatch
+    this[kOriginalClose] = this.close.bind(this)
+
+    this.dispatch = buildMockDispatch.call(this)
+    this.close = this[kClose]
+  }
+
+  get [Symbols.kConnected] () {
+    return this[kConnected]
+  }
+
+  /**
+   * Sets up the base interceptor for mocking replies from undici.
+   */
+  intercept (opts) {
+    return new MockInterceptor(opts, this[kDispatches])
+  }
+
+  async [kClose] () {
+    await promisify(this[kOriginalClose])()
+    this[kConnected] = 0
+    this[kMockAgent][Symbols.kClients].delete(this[kOrigin])
+  }
+}
+
+module.exports = MockPool