Prevent infinite loop in Command::Base#step & add configurable wait time (#217)

Author: zzaakiirr

.rubocop.yml

@@ -20,3 +20,7 @@ RSpec/ExampleLength:
 
 RSpec/MultipleExpectations:
   Enabled: false
+
+RSpec/NestedGroups:
+  Enabled: true
+  Max: 5

lib/command/base.rb

@@ -474,21 +474,19 @@ def step_finish(success, abort_on_error: true)
       end
     end
 
-    def step(message, abort_on_error: true, retry_on_failure: false) # rubocop:disable Metrics/MethodLength
+    def step(message, abort_on_error: true, retry_on_failure: false, max_retry_count: 1000, wait: 1, &block) # rubocop:disable Metrics/MethodLength
       progress.print("#{message}...")
 
       Shell.use_tmp_stderr do
         success = false
 
         begin
-          if retry_on_failure
-            until (success = yield)
-              progress.print(".")
-              Kernel.sleep(1)
+          success =
+            if retry_on_failure
+              with_retry(max_retry_count: max_retry_count, wait: wait, &block)
+            else
+              yield
             end
-          else
-            success = yield
-          end
         rescue RuntimeError => e
           Shell.write_to_tmp_stderr(e.message)
         end
@@ -497,6 +495,22 @@ def step(message, abort_on_error: true, retry_on_failure: false) # rubocop:disab
       end
     end
 
+    def with_retry(max_retry_count:, wait:)
+      retry_count = 0
+      success = false
+
+      while !success && retry_count <= max_retry_count
+        success = yield
+        break if success
+
+        progress.print(".")
+        Kernel.sleep(wait)
+        retry_count += 1
+      end
+
+      success
+    end
+
     def cp
       @cp ||= Controlplane.new(config)
     end

spec/command/base_spec.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+describe Command::Base do
+  let(:config) { instance_double(Command::Config) }
+  let(:command) { described_class.new(config) }
+
+  around do |example|
+    suppress_output { example.run }
+  end
+
+  describe "#step" do
+    let(:message) { "test message" }
+    let(:common_options) { { abort_on_error: false } }
+
+    context "with retry_on_failure: true" do
+      let(:options) { common_options.merge(retry_on_failure: true, wait: 0) }
+
+      it "retries block until success" do
+        run_count = 0
+
+        command.step(message, **options) do
+          run_count += 1
+          true if run_count == 3
+        end
+
+        expect(run_count).to eq(3)
+      end
+
+      it "does not exceed default max_retry_count" do
+        run_count = 0
+
+        command.step(message, **options) do
+          run_count += 1
+          false
+        end
+
+        expect(run_count).to eq(1001)
+      end
+
+      context "with max_retry_count option" do # rubocop:disable RSpec/MultipleMemoizedHelpers
+        let(:options_with_max_retry_count) { common_options.merge(retry_on_failure: true, wait: 0, max_retry_count: 1) }
+
+        it "retries block specified times" do
+          run_count = 0
+
+          command.step(message, **options_with_max_retry_count) do
+            run_count += 1
+            false
+          end
+
+          expect(run_count).to eq(2)
+        end
+      end
+    end
+
+    context "with retry_on_failure: false" do
+      let(:options) { common_options.merge(retry_on_failure: false) }
+
+      it "does not retry block" do
+        run_count = 0
+
+        command.step(message, **options) do
+          run_count += 1
+          false
+        end
+
+        expect(run_count).to eq(1)
+      end
+    end
+  end
+end

spec/support/command_helpers.rb

@@ -207,6 +207,16 @@ def spawn_cpflow_command(*args, stty_rows: nil, stty_cols: nil, wait_for_process
     end
   end
 
+  def suppress_output
+    original_stderr = replace_stderr
+    original_stdout = replace_stdout
+
+    yield
+  ensure
+    restore_stderr(original_stderr)
+    restore_stdout(original_stdout)
+  end
+
   def replace_stderr
     original_stderr = $stderr
     $stderr = Tempfile.create