feat: add factory option to BalancedPool

nodejs · Feb 16, 2022 · ec67852 · ec67852
1 parent 5627f87
commit ec67852
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/lib/balanced-pool.js b/lib/balanced-pool.js
@@ -1,7 +1,8 @@
 'use strict'
 
 const {
-  BalancedPoolMissingUpstreamError
+  BalancedPoolMissingUpstreamError,
+  InvalidArgumentError
 } = require('./core/errors')
 const {
   PoolBase,
@@ -13,11 +14,16 @@ const {
 } = require('./pool-base')
 const Pool = require('./pool')
 const { kUrl } = require('./core/symbols')
+const kFactory = Symbol('factory')
 
 const kOptions = Symbol('options')
 
+function defaultFactory (origin, opts) {
+  return new Pool(origin, opts);
+}
+
 class BalancedPool extends PoolBase {
-  constructor (upstreams = [], opts = {}) {
+  constructor (upstreams = [], { factory = defaultFactory, ...opts } = {}) {
     super()
 
     this[kOptions] = opts
@@ -26,6 +32,12 @@ class BalancedPool extends PoolBase {
       upstreams = [upstreams]
     }
 
+    if (typeof factory !== 'function') {
+      throw new InvalidArgumentError('factory must be a function.')
+    }
+
+    this[kFactory] = factory
+
     for (const upstream of upstreams) {
       this.addUpstream(upstream)
     }
@@ -40,7 +52,7 @@ class BalancedPool extends PoolBase {
       return this
     }
 
-    this[kAddClient](new Pool(upstream, Object.assign({}, this[kOptions])))
+    this[kAddClient](this[kFactory](upstream, Object.assign({}, this[kOptions])))
 
     return this
   }

diff --git a/test/balanced-pool.js b/test/balanced-pool.js
@@ -163,3 +163,14 @@ test('busy', (t) => {
     }
   })
 })
+
+test('invalid options throws', (t) => {
+  t.plan(2)
+
+  try {
+    new BalancedPool(null, { factory: '' }) // eslint-disable-line
+  } catch (err) {
+    t.type(err, errors.InvalidArgumentError)
+    t.equal(err.message, 'factory must be a function.')
+  }
+})