allow procs as command map values

CHANGELOG.md

@@ -32,6 +32,9 @@ appear at the top.
     and have a cleaner internal API. You can still completely disable the pool
     by setting `SSHKit::Backend::Netssh.pool.idle_timeout = 0`.
     @mattbrictson @byroot [PR #328](https://github.com/capistrano/sshkit/pull/328)
+  * Allow command map entries (`SSHKit::CommandMap#[]`) to be Procs
+    [PR #310]((https://github.com/capistrano/sshkit/pull/310)
+    @mikz
 
 ### Bug fixes
 

lib/sshkit/command_map.rb

@@ -31,18 +31,20 @@ def [](command)
       end
     end
 
+    TO_VALUE = ->(obj) { obj.respond_to?(:call) ? obj.call : obj }
+
     def initialize(value = nil)
       @map = CommandHash.new(value || defaults)
     end
 
     def [](command)
       if prefix[command].any?
-        prefixes = prefix[command].map{ |prefix| prefix.respond_to?(:call) ? prefix.call : prefix }
+        prefixes = prefix[command].map(&TO_VALUE)
         prefixes = prefixes.join(" ")
 
         "#{prefixes} #{command}"
       else
-        @map[command]
+        TO_VALUE.(@map[command])
       end
     end
 

test/unit/test_command_map.rb

@@ -16,6 +16,15 @@ def test_setter
       assert_equal map[:rake], "/usr/local/rbenv/shims/rake"
     end
 
+    def test_setter_procs
+      map = CommandMap.new
+      i = 0
+      map[:rake] = -> { i += 1; "/usr/local/rbenv/shims/rake#{i}" }
+
+      assert_equal map[:rake], "/usr/local/rbenv/shims/rake1"
+      assert_equal map[:rake], "/usr/local/rbenv/shims/rake2"
+    end
+
     def test_prefix
       map = CommandMap.new
       map.prefix[:rake].push("/home/vagrant/.rbenv/bin/rbenv exec")