From ec678523567a2822f30f2a99a9f9b6812f27583f Mon Sep 17 00:00:00 2001
From: Arpan KC <nipeshkc7@gmail.com>
Date: Wed, 16 Feb 2022 17:06:37 +1100
Subject: [PATCH] feat: add factory option to BalancedPool

---
 lib/balanced-pool.js  | 18 +++++++++++++++---
 test/balanced-pool.js | 11 +++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/lib/balanced-pool.js b/lib/balanced-pool.js
index 4ad0bf48602..641c813b152 100644
--- 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
index 473522ac9f2..1fd6ff82618 100644
--- 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.')
+  }
+})
\ No newline at end of file