diff --git a/lib/kafo/puppet_command.rb b/lib/kafo/puppet_command.rb
index 7e1f4b1b..e8de2aac 100644
--- a/lib/kafo/puppet_command.rb
+++ b/lib/kafo/puppet_command.rb
@@ -50,11 +50,21 @@ def self.aio_puppet?
       false
     end
 
+    # Format a command for use in Open3 that properly filters the environment
+    # for use with Puppet
+    #
+    # @param command [String, Array]
+    #   The command to execute. A string implies it'll run with a shell while
+    #   an array runs without. That removes the need to escape input and is
+    #   recommended.
+    # @return [Array]
+    #   A command for use in Open3
     def self.format_command(command)
+      cmd = command.is_a?(Array) ? command : [command]
       if aio_puppet?
-        [clean_env_vars, command, :unsetenv_others => true]
+        [clean_env_vars] + cmd + [:unsetenv_others => true]
       else
-        [::ENV, command, :unsetenv_others => false]
+        [::ENV] + cmd + [:unsetenv_others => false]
       end
     end
 
diff --git a/test/kafo/puppet_command_test.rb b/test/kafo/puppet_command_test.rb
index fe5869ee..05f5225d 100644
--- a/test/kafo/puppet_command_test.rb
+++ b/test/kafo/puppet_command_test.rb
@@ -57,7 +57,7 @@ module Kafo
       describe "with 'puppet' in PATH" do
         specify do
           ::ENV.stub(:[], '/usr/bin:/usr/local/bin') do
-            File.stub(:executable?, Proc.new { |path| path == '/usr/local/bin/puppet' }) do
+            File.stub(:executable?, proc { |path| path == '/usr/local/bin/puppet' }) do
               _(pc).must_equal '/usr/local/bin/puppet'
             end
           end
@@ -67,7 +67,7 @@ module Kafo
       describe "with AIO 'puppet' only" do
         specify do
           ::ENV.stub(:[], '/usr/bin:/usr/local/bin') do
-            File.stub(:executable?, Proc.new { |path| path == '/opt/puppetlabs/bin/puppet' }) do
+            File.stub(:executable?, proc { |path| path == '/opt/puppetlabs/bin/puppet' }) do
               _(pc).must_equal '/opt/puppetlabs/bin/puppet'
             end
           end
@@ -90,7 +90,7 @@ module Kafo
         let(:puppet_command) { '/usr/bin/puppet' }
 
         specify 'as a real file' do
-          File.stub(:realpath, ->(path) { path }) do
+          File.stub(:realpath, proc { |path| path }) do
             refute subject
           end
         end
@@ -102,7 +102,7 @@ module Kafo
         end
 
         specify 'as a broken symlink' do
-          File.stub(:realpath, ->(path) { raise Errno::ENOENT, 'No such file or directory' }) do
+          File.stub(:realpath, proc { |_path| raise Errno::ENOENT, 'No such file or directory' }) do
             refute subject
           end
         end
@@ -112,11 +112,49 @@ module Kafo
         let(:puppet_command) { 'puppet' }
 
         specify 'non-existant' do
-          File.stub(:realpath, ->(path) { raise Errno::ENOENT, 'No such file or directory' }) do
+          File.stub(:realpath, proc { |_path| raise Errno::ENOENT, 'No such file or directory' }) do
             refute subject
           end
         end
       end
     end
+
+    describe '.format_command' do
+      describe 'with Puppet AIO' do
+        specify 'with a string' do
+          PuppetCommand.stub(:aio_puppet?, true) do
+            PuppetCommand.stub(:clean_env_vars, {'FOO' => 'bar'}) do
+              expected = [{'FOO' => 'bar'}, 'echo hello world', { :unsetenv_others => true }]
+              assert_equal(expected, PuppetCommand.format_command('echo hello world'))
+            end
+          end
+        end
+
+        specify 'with an array' do
+          PuppetCommand.stub(:aio_puppet?, true) do
+            PuppetCommand.stub(:clean_env_vars, {'FOO' => 'bar'}) do
+              expected = [{'FOO' => 'bar'}, 'echo', 'hello', 'world', { :unsetenv_others => true }]
+              assert_equal(expected, PuppetCommand.format_command(['echo', 'hello', 'world']))
+            end
+          end
+        end
+      end
+
+      describe 'with regular Puppet' do
+        specify 'with a string' do
+          PuppetCommand.stub(:aio_puppet?, false) do
+            expected = [::ENV, 'echo hello world', { :unsetenv_others => false }]
+            assert_equal(expected, PuppetCommand.format_command('echo hello world'))
+          end
+        end
+
+        specify 'with an array' do
+          PuppetCommand.stub(:aio_puppet?, false) do
+            expected = [::ENV, 'echo', 'hello', 'world', { :unsetenv_others => false }]
+            assert_equal(expected, PuppetCommand.format_command(['echo', 'hello', 'world']))
+          end
+        end
+      end
+    end
   end
 end