Allow passing extra environment variables to PuppetCommand.format_com…

…mand

lib/kafo/puppet_command.rb

@@ -59,12 +59,12 @@ def self.aio_puppet?
     #   recommended.
     # @return [Array]
     #   A command for use in Open3
-    def self.format_command(command)
+    def self.format_command(command, extra_env = {})
       cmd = command.is_a?(Array) ? command : [command]
       if aio_puppet?
-        [clean_env_vars] + cmd + [:unsetenv_others => true]
+        [clean_env_vars.merge(extra_env)] + cmd + [:unsetenv_others => true]
       else
-        [::ENV] + cmd + [:unsetenv_others => false]
+        [::ENV.to_h.merge(extra_env)] + cmd + [:unsetenv_others => false]
       end
     end
 

test/kafo/puppet_command_test.rb

@@ -138,6 +138,15 @@ module Kafo
             end
           end
         end
+
+        specify 'with extra env' do
+          PuppetCommand.stub(:aio_puppet?, true) do
+            PuppetCommand.stub(:clean_env_vars, {'FOO' => 'bar'}) do
+              expected = [{'FOO' => 'bar', 'EXTRA' => 'cool'}, 'echo hello world', { :unsetenv_others => true }]
+              assert_equal(expected, PuppetCommand.format_command(['echo hello world'], { 'EXTRA' => 'cool' }))
+            end
+          end
+        end
       end
 
       describe 'with regular Puppet' do
@@ -154,6 +163,15 @@ module Kafo
             assert_equal(expected, PuppetCommand.format_command(['echo', 'hello', 'world']))
           end
         end
+
+        specify 'with extra env' do
+          PuppetCommand.stub(:aio_puppet?, false) do
+            expected_env = ::ENV.to_h.merge({ 'EXTRA' => 'cool' })
+            expected = [expected_env, 'echo hello world', { :unsetenv_others => false }]
+            assert_equal(expected, PuppetCommand.format_command('echo hello world', { 'EXTRA' => 'cool' }))
+          end
+        end
+
       end
     end
   end