From 278de3e30f94cb49c58eb10391d4d79851e3d572 Mon Sep 17 00:00:00 2001
From: Michal Cichra <michal@o2h.cz>
Date: Mon, 28 Dec 2015 15:24:11 +0100
Subject: [PATCH] allow procs as command map values

so this example works properly:

  SSHKit.config.command_map[:bundle] = -> { "#{fetch(:ruby_cmd)} /usr/bin/local/bundle" }
---
 CHANGELOG.md                  | 3 +++
 lib/sshkit/command_map.rb     | 6 ++++--
 test/unit/test_command_map.rb | 9 +++++++++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index dd40af34..0bbaa877 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,9 @@ appear at the top.
   * display more accurate string for commands with spaces being output in `Formatter::Pretty`
     [PR #304](https://github.com/capistrano/sshkit/pull/304)
     @steved
+  * allow command map entries to be Procs
+    [PR #310]((https://github.com/capistrano/sshkit/pull/310)
+    @mikz
 
 ## 1.8.1
 
diff --git a/lib/sshkit/command_map.rb b/lib/sshkit/command_map.rb
index 13dedcd8..88064ad0 100644
--- a/lib/sshkit/command_map.rb
+++ b/lib/sshkit/command_map.rb
@@ -33,18 +33,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
 
diff --git a/test/unit/test_command_map.rb b/test/unit/test_command_map.rb
index eb0e46d0..90bddd89 100644
--- a/test/unit/test_command_map.rb
+++ b/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")