From 453cd9e1af372d3839a72da8817d70849a3fcd9f Mon Sep 17 00:00:00 2001
From: spalger <spalger@users.noreply.github.com>
Date: Thu, 4 Feb 2016 12:12:25 -0700
Subject: [PATCH] [transport] expose a method for updating the host list

---
 src/lib/connection.js                  |  2 --
 src/lib/transport.js                   | 34 +++++++++++++----------
 test/unit/specs/connection_abstract.js |  9 -------
 test/unit/specs/transport.js           | 37 ++++++++++++++++++++++++++
 4 files changed, 57 insertions(+), 25 deletions(-)

diff --git a/src/lib/connection.js b/src/lib/connection.js
index 3e16f3e12..93c8c38a6 100644
--- a/src/lib/connection.js
+++ b/src/lib/connection.js
@@ -22,8 +22,6 @@ function ConnectionAbstract(host, config) {
     throw new TypeError('Missing host');
   } else if (host instanceof Host) {
     this.host = host;
-  } else if (typeof host === "string"){
-    this.host = new Host(host);
   } else {
     throw new TypeError('Invalid host');
   }
diff --git a/src/lib/transport.js b/src/lib/transport.js
index 71aa7bb50..f8adb6668 100644
--- a/src/lib/transport.js
+++ b/src/lib/transport.js
@@ -61,15 +61,11 @@ function Transport(config) {
         'or an array of host config objects.');
     }
 
-    var hosts = _.map(hostsConfig, function (conf) {
-      return (conf instanceof Host) ? conf : new Host(conf, self._config);
-    });
-
     if (randomizeHosts) {
-      hosts = _.shuffle(hosts);
+      hostsConfig = _.shuffle(hostsConfig);
     }
 
-    self.connectionPool.setHosts(hosts);
+    self.setHosts(hostsConfig);
   }
 
   if (config.hasOwnProperty('sniffedNodesProtocol')) {
@@ -390,10 +386,9 @@ Transport.prototype._timeout = function (cb, delay) {
  * @param  {Function} cb - Function to call back once complete
  */
 Transport.prototype.sniff = function (cb) {
-  var connectionPool = this.connectionPool;
+  var self = this;
   var nodesToHostCallback = this.nodesToHostCallback;
   var log = this.log;
-  var globalConfig = this._config;
   var sniffedNodesProtocol = this.sniffedNodesProtocol;
 
   // make cb a function if it isn't
@@ -414,18 +409,29 @@ Transport.prototype.sniff = function (cb) {
         return;
       }
 
-      connectionPool.setHosts(_.map(hostsConfigs, function (hostConfig) {
-        if (sniffedNodesProtocol) {
-          hostConfig.protocol = sniffedNodesProtocol;
-        }
+      _.forEach(hostsConfigs, function (hostConfig) {
+        if (sniffedNodesProtocol) hostConfig.protocol = sniffedNodesProtocol;
+      });
 
-        return new Host(hostConfig, globalConfig);
-      }));
+      self.setHosts(hostsConfigs);
     }
     cb(err, resp, status);
   });
 };
 
+/**
+ * Set the host list that the transport should use.
+ *
+ * @param {Array<HostConfig>} hostsConfigs - an array of Hosts, or configuration objects
+ *                                         that will be used to create Host objects.
+ */
+Transport.prototype.setHosts = function (hostsConfigs) {
+  var globalConfig = this._config;
+  this.connectionPool.setHosts(_.map(hostsConfigs, function (conf) {
+    return (conf instanceof Host) ? conf : new Host(conf, globalConfig);
+  }));
+};
+
 /**
  * Close the Transport, which closes the logs and connection pool
  * @return {[type]} [description]
diff --git a/test/unit/specs/connection_abstract.js b/test/unit/specs/connection_abstract.js
index 786436c68..a1e8ec88b 100644
--- a/test/unit/specs/connection_abstract.js
+++ b/test/unit/specs/connection_abstract.js
@@ -15,15 +15,6 @@ describe('Connection Abstract', function () {
     expect(conn.host).to.be(host);
   });
 
-  it('constructs with string host using Host constructor', function(){
-    var host = 'localhost';
-    var port = 9200;
-    var hostUrl = host+':'+port;
-    var conn = new ConnectionAbstract(hostUrl);
-    expect(conn.host.host).to.be(host);
-    expect(conn.host.port).to.be(port);
-  });
-
   it('requires a valid host', function () {
     expect(function () {
       new ConnectionAbstract();
diff --git a/test/unit/specs/transport.js b/test/unit/specs/transport.js
index 46b0625b1..39c15fca6 100644
--- a/test/unit/specs/transport.js
+++ b/test/unit/specs/transport.js
@@ -807,6 +807,43 @@ describe('Transport Class', function () {
     });
   });
 
+  describe('#setHosts', function () {
+    it('accepts strings, host objects, and host configs', function () {
+
+      var trans = new Transport({ suggestCompression: true });
+      stub(trans.connectionPool, 'setHosts');
+
+      trans.setHosts([
+        { host: 'first.server', port: 9200 },
+        'http://second.server:9200',
+        new Host('http://third.server:9200')
+      ]);
+
+      sinon.assert.calledOnce(trans.connectionPool.setHosts);
+      var host, hosts = trans.connectionPool.setHosts.firstCall.args[0];
+
+      expect(hosts).to.have.length(3);
+
+      host = hosts.shift();
+      expect(host).to.be.a(Host);
+      expect(host.host).to.eql('first.server');
+      expect(host.port).to.eql(9200);
+      expect(host.suggestCompression).to.be(true);
+
+      host = hosts.shift();
+      expect(host).to.be.a(Host);
+      expect(host.host).to.eql('second.server');
+      expect(host.port).to.eql(9200);
+      expect(host.suggestCompression).to.be(true);
+
+      host = hosts.shift();
+      expect(host).to.be.a(Host);
+      expect(host.host).to.eql('third.server');
+      expect(host.port).to.eql(9200);
+      expect(host.suggestCompression).to.be(false);
+    });
+  });
+
   describe('#close', function () {
     it('proxies the call to it\'s log and connection pool', function () {
       var tran = new Transport();