diff --git a/lib/kafo/puppet_command.rb b/lib/kafo/puppet_command.rb index e8de2aac..0e5edbf6 100644 --- a/lib/kafo/puppet_command.rb +++ b/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 diff --git a/test/kafo/puppet_command_test.rb b/test/kafo/puppet_command_test.rb index 05f5225d..1fd61940 100644 --- a/test/kafo/puppet_command_test.rb +++ b/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