Make tcp-keepalive configurable (brianc#1058)

brianc · web-flow · commit 0aa62f2854fa · 2016-06-21T21:38:10.000-05:00
diff --git a/lib/client.js b/lib/client.js
@@ -25,7 +25,8 @@ var Client = function(config) {
 
   this.connection = c.connection || new Connection({
     stream: c.stream,
-    ssl: this.connectionParameters.ssl
+    ssl: this.connectionParameters.ssl,
+    keepAlive: c.keepAlive || false
   });
   this.queryQueue = [];
   this.binary = c.binary || defaults.binary;
diff --git a/lib/connection.js b/lib/connection.js
@@ -11,6 +11,7 @@ var Connection = function(config) {
   EventEmitter.call(this);
   config = config || {};
   this.stream = config.stream || new net.Stream();
+  this._keepAlive = config.keepAlive;
   this.lastBuffer = false;
   this.lastOffset = 0;
   this.buffer = null;
@@ -47,8 +48,10 @@ Connection.prototype.connect = function(port, host) {
   var self = this;
 
   this.stream.on('connect', function() {
+    if (self._keepAlive) {
+      self.stream.setKeepAlive(true);
+    }
     self.emit('connect');
-    self.stream.setKeepAlive(true);
   });
 
   this.stream.on('error', function(error) {
diff --git a/test/unit/connection/startup-tests.js b/test/unit/connection/startup-tests.js
@@ -7,13 +7,18 @@ test('connection can take existing stream', function() {
 });
 
 test('using closed stream', function() {
-  var stream = new MemoryStream();
-  stream.readyState = 'closed';
-  stream.connect = function(port, host) {
-    this.connectCalled = true;
-    this.port = port;
-    this.host = host;
-  }
+  var makeStream = function() {
+    var stream = new MemoryStream();
+    stream.readyState = 'closed';
+    stream.connect = function(port, host) {
+      this.connectCalled = true;
+      this.port = port;
+      this.host = host;
+    }
+    return stream;
+  };
+
+  var stream = makeStream();
 
   var con = new Connection({stream: stream});
 
@@ -46,14 +51,20 @@ test('using closed stream', function() {
 
   test('after stream emits connected event init TCP-keepalive', function() {
 
+    var stream = makeStream();
+    var con = new Connection({ stream: stream, keepAlive: true });
+    con.connect(123, 'test');
+
     var res = false;
 
     stream.setKeepAlive = function(bit) {
       res = bit;
     };
 
     assert.ok(stream.emit('connect'));
-    assert.equal(res, true);
+    setTimeout(function() {
+      assert.equal(res, true);
+    })
   });
 });
 
