From 43d6aa23dd9bfd3041e03496a8a9ab9d8f2ee0fc Mon Sep 17 00:00:00 2001
From: Hiroshi Shirosaki <h.shirosaki@gmail.com>
Date: Wed, 8 Mar 2017 15:26:58 +0900
Subject: [PATCH] Fix test for actual bug and fix updating key

`Host#netssh_options` creates new Hash from `ssh_options` when
calling `with`. So new Hash would be required.
:password_prompt and :logger options are added in net-ssh
`Net::SSH.assign_defaults`.
The key in Cache object should be updated for comparing next time.
---
 lib/sshkit/backends/connection_pool.rb       | 10 +++++++---
 lib/sshkit/backends/connection_pool/cache.rb |  2 +-
 test/unit/backends/test_connection_pool.rb   | 13 +++++++++----
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/lib/sshkit/backends/connection_pool.rb b/lib/sshkit/backends/connection_pool.rb
index b0c769c9..4f03352b 100644
--- a/lib/sshkit/backends/connection_pool.rb
+++ b/lib/sshkit/backends/connection_pool.rb
@@ -54,6 +54,8 @@ def initialize(idle_timeout=30)
   # invoking the `connection_factory` proc with the given `args`. The arguments
   # are used to construct a key used for caching.
   def with(connection_factory, *args)
+    options = args.last
+    args[args.size - 1] = options.dup if options.is_a? Hash
     cache = find_cache(args)
     conn = cache.pop || begin
       connection_factory.call(*args)
@@ -63,7 +65,7 @@ def with(connection_factory, *args)
     cache.push(conn) unless conn.nil?
     # Sometimes the args mutate as a result of opening a connection. In this
     # case we need to update the cache key to match the new args.
-    update_key_if_args_changed(cache, args)
+    update_key_if_args_changed(cache, args, options)
   end
 
   # Immediately remove all cached connections, without closing them. This only
@@ -116,12 +118,14 @@ def thread_safe_find_or_create_cache(key)
   end
 
   # Update cache key with changed args to prevent cache miss
-  def update_key_if_args_changed(cache, args)
+  def update_key_if_args_changed(cache, args, options)
+    args[args.size - 1] = options if options.is_a? Hash
     new_key = cache_key_for_connection_args(args)
     return if cache.same_key?(new_key)
 
     caches.synchronize do
-      caches[new_key] = caches.delete(cache.key)
+      cache = caches[new_key] = caches.delete(cache.key)
+      cache.key = new_key
     end
   end
 
diff --git a/lib/sshkit/backends/connection_pool/cache.rb b/lib/sshkit/backends/connection_pool/cache.rb
index ffe0adcc..8be4757e 100644
--- a/lib/sshkit/backends/connection_pool/cache.rb
+++ b/lib/sshkit/backends/connection_pool/cache.rb
@@ -1,7 +1,7 @@
 # A Cache holds connections for a given key. Each connection is stored along
 # with an expiration time so that its idle duration can be measured.
 class SSHKit::Backend::ConnectionPool::Cache
-  attr_reader :key
+  attr_accessor :key
 
   def initialize(key, idle_timeout, closer)
     @key = key
diff --git a/test/unit/backends/test_connection_pool.rb b/test/unit/backends/test_connection_pool.rb
index c8a361a9..4cd9246e 100644
--- a/test/unit/backends/test_connection_pool.rb
+++ b/test/unit/backends/test_connection_pool.rb
@@ -136,10 +136,15 @@ def test_close_connections
       end
 
       def test_connections_with_changed_args_is_reused
-        options = { known_hosts: "foo" }
-        connect_change_options = ->(*args) { args.last[:known_hosts] = "bar"; Object.new }
-        conn1 = pool.with(connect_change_options, "arg", options) { |c| c }
-        conn2 = pool.with(connect_change_options, "arg", options) { |c| c }
+        options = { known_hosts: ["foo"] }
+        connect_change_options = ->(*args) {
+          args.last[:known_hosts] << "bar"
+          args.last[:password_prompt] ||= Object.new
+          args.last[:logger] ||= Object.new
+          Object.new
+        }
+        conn1 = pool.with(connect_change_options, "arg", {}.merge(options)) { |c| c }
+        conn2 = pool.with(connect_change_options, "arg", {}.merge(options)) { |c| c }
 
         assert_equal conn1, conn2
       end