From 15baea11102cf7b0322f5113296aba66c31864c2 Mon Sep 17 00:00:00 2001
From: Zakir Dzhamaliddinov <zakir.dzhamaliddinov@gmail.com>
Date: Mon, 22 Jul 2024 17:42:31 +0300
Subject: [PATCH] Add max_retry_count & wait options to Command::Base#step

---
 .rubocop.yml                    |  4 ++
 lib/command/base.rb             | 12 ++++--
 spec/command/base_spec.rb       | 73 +++++++++++++++++++++++++++++++++
 spec/support/command_helpers.rb | 10 +++++
 4 files changed, 96 insertions(+), 3 deletions(-)
 create mode 100644 spec/command/base_spec.rb

diff --git a/.rubocop.yml b/.rubocop.yml
index 29d37195..f1a81d5c 100644
--- a/.rubocop.yml
+++ b/.rubocop.yml
@@ -20,3 +20,7 @@ RSpec/ExampleLength:
 
 RSpec/MultipleExpectations:
   Enabled: false
+
+RSpec/NestedGroups:
+  Enabled: true
+  Max: 5
diff --git a/lib/command/base.rb b/lib/command/base.rb
index b0c645e3..eac48bd9 100644
--- a/lib/command/base.rb
+++ b/lib/command/base.rb
@@ -483,7 +483,7 @@ def step_finish(success)
       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: 5, wait: 1) # rubocop:disable Metrics/MethodLength
       progress.print("#{message}...")
 
       Shell.use_tmp_stderr do
@@ -491,9 +491,15 @@ def step(message, abort_on_error: true, retry_on_failure: false) # rubocop:disab
 
         begin
           if retry_on_failure
-            until (success = yield)
+            run_count = 0
+            while !success && run_count < max_retry_count
+              success = yield
+              break if success
+
               progress.print(".")
-              Kernel.sleep(1)
+              Kernel.sleep(wait)
+
+              run_count += 1
             end
           else
             success = yield
diff --git a/spec/command/base_spec.rb b/spec/command/base_spec.rb
new file mode 100644
index 00000000..8a58f029
--- /dev/null
+++ b/spec/command/base_spec.rb
@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+describe Command::Base do
+  subject(:command) { described_class.new(config) }
+
+  let(:config) { instance_double(Command::Config) }
+
+  around do |example|
+    suppress_output { example.run }
+  end
+
+  describe "#step" do
+    let(:message) { "test message" }
+
+    context "with retry_on_failure: true" do
+      let(:options) { { retry_on_failure: true, wait: 0 } }
+
+      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 5
+      end
+
+      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
+
+      context "with max_retry_count option" do
+        let(:options) { super().merge(max_retry_count: 1) }
+
+        it "retries block specified times" do
+          run_count = 0
+
+          command.step(message, **options) do
+            run_count += 1
+            false
+          end
+
+          expect(run_count).to eq 1
+        end
+      end
+    end
+
+    context "with retry_on_failure: false" do
+      let(:options) { { 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
diff --git a/spec/support/command_helpers.rb b/spec/support/command_helpers.rb
index 6e275083..d4878964 100644
--- a/spec/support/command_helpers.rb
+++ b/spec/support/command_helpers.rb
@@ -203,6 +203,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