Add timeout support to Mixlib::ShellOut based local runners

lib/train/transports/local.rb

@@ -86,9 +86,9 @@ def force_runner(command_runner)
         end
       end
 
-      def run_command_via_connection(cmd, &_data_handler)
+      def run_command_via_connection(cmd, opts, &_data_handler)
         # Use the runner if it is available
-        return @runner.run_command(cmd) if defined?(@runner)
+        return @runner.run_command(cmd, opts) if defined?(@runner)
 
         # If we don't have a runner, such as at the beginning of setting up the
         # transport and performing the first few steps of OS detection, fall
@@ -115,13 +115,18 @@ def initialize(connection, options)
           @cmd_wrapper = Local::CommandWrapper.load(connection, options)
         end
 
-        def run_command(cmd)
+        def run_command(cmd, opts = {})
           if defined?(@cmd_wrapper) && !@cmd_wrapper.nil?
             cmd = @cmd_wrapper.run(cmd)
           end
 
           res = Mixlib::ShellOut.new(cmd)
-          res.run_command
+          res.timeout = opts[:timeout]
+          begin
+            res.run_command
+          rescue Mixlib::ShellOut::CommandTimeout
+            raise Train::CommandTimeoutReached
+          end
           Local::CommandResult.new(res.stdout, res.stderr, res.exitstatus)
         end
 
@@ -138,7 +143,7 @@ def initialize(powershell_cmd = "powershell")
           @powershell_cmd = powershell_cmd
         end
 
-        def run_command(script)
+        def run_command(script, opts)
           # Prevent progress stream from leaking into stderr
           script = "$ProgressPreference='SilentlyContinue';" + script
 
@@ -149,7 +154,12 @@ def run_command(script)
           cmd = "#{@powershell_cmd} -NoProfile -EncodedCommand #{base64_script}"
 
           res = Mixlib::ShellOut.new(cmd)
-          res.run_command
+          res.timeout = opts[:timeout]
+          begin
+            res.run_command
+          rescue Mixlib::ShellOut::CommandTimeout
+            raise Train::CommandTimeoutReached
+          end
           Local::CommandResult.new(res.stdout, res.stderr, res.exitstatus)
         end
 
@@ -176,9 +186,10 @@ def initialize(powershell_cmd = "powershell")
         #         A command that succeeds without setting an exit code will have exitstatus 0
         #         A command that exits with an exit code will have that value as exitstatus
         #         A command that fails (e.g. throws exception) before setting an exit code will have exitstatus 1
-        def run_command(cmd)
+        def run_command(cmd, _opts)
           script = "$ProgressPreference='SilentlyContinue';" + cmd
           encoded_script = Base64.strict_encode64(script)
+          # TODO: no way to safely implement timeouts here.
           @pipe.puts(encoded_script)
           @pipe.flush
           res = OpenStruct.new(JSON.parse(Base64.decode64(@pipe.readline)))

test/unit/transports/local_test.rb

@@ -113,6 +113,7 @@ def initialize(user_opts = {})
 
     def mock_run_cmd(cmd, &block)
       cmd_runner.expect :run_command, nil
+      cmd_runner.expect :timeout=, nil, [nil]
       Mixlib::ShellOut.stub :new, cmd_runner do |*args|
         yield
       end
@@ -164,7 +165,7 @@ def mock_run_cmd(cmd, &block)
         .expects(:new)
         .never
 
-      runner.expects(:run_command).with("not actually executed")
+      runner.expects(:run_command).with("not actually executed", {})
       connection.run_command("not actually executed")
     end
 
@@ -177,7 +178,7 @@ def mock_run_cmd(cmd, &block)
         .expects(:new)
         .returns(runner)
 
-      runner.expects(:run_command).with("not actually executed")
+      runner.expects(:run_command).with("not actually executed", {})
       connection.run_command("not actually executed")
     end
   end