1 files changed, 43 insertions, 31 deletions
diff --git a/qa/qa/support/retrier.rb b/qa/qa/support/retrier.rb
index 3b02cb4855b..7b548e95453 100644
--- a/qa/qa/support/retrier.rb
+++ b/qa/qa/support/retrier.rb
@@ -3,49 +3,61 @@
 module QA
   module Support
     module Retrier
+      extend Repeater
+
       module_function
 
       def retry_on_exception(max_attempts: 3, reload_page: nil, sleep_interval: 0.5)
-        QA::Runtime::Logger.debug("with retry_on_exception: max_attempts #{max_attempts}; sleep_interval #{sleep_interval}")
-
-        attempts = 0
+        QA::Runtime::Logger.debug(
+          <<~MSG.tr("\n", ' ')
+            with retry_on_exception: max_attempts: #{max_attempts};
+            reload_page: #{reload_page};
+            sleep_interval: #{sleep_interval}
+        MSG
+        )
 
-        begin
-          QA::Runtime::Logger.debug("Attempt number #{attempts + 1}")
-          yield
-        rescue StandardError, RSpec::Expectations::ExpectationNotMetError
-          sleep sleep_interval
-          reload_page.refresh if reload_page
-          attempts += 1
+        result = nil
+        repeat_until(
+          max_attempts: max_attempts,
+          reload_page: reload_page,
+          sleep_interval: sleep_interval,
+          retry_on_exception: true
+        ) do
+          result = yield
 
-          retry if attempts < max_attempts
-          QA::Runtime::Logger.debug("Raising exception after #{max_attempts} attempts")
-          raise
+          # This method doesn't care what the return value of the block is.
+          # We set it to `true` so that it doesn't repeat if there's no exception
+          true
         end
-      end
-
-      def retry_until(max_attempts: 3, reload_page: nil, sleep_interval: 0, exit_on_failure: false)
-        QA::Runtime::Logger.debug("with retry_until: max_attempts #{max_attempts}; sleep_interval #{sleep_interval}; reload_page:#{reload_page}")
-        attempts = 0
+        QA::Runtime::Logger.debug("ended retry_on_exception")
 
-        while attempts < max_attempts
-          QA::Runtime::Logger.debug("Attempt number #{attempts + 1}")
-          result = yield
-          return result if result
+        result
+      end
 
-          sleep sleep_interval
+      def retry_until(max_attempts: nil, max_duration: nil, reload_page: nil, sleep_interval: 0, raise_on_failure: false, retry_on_exception: false)
+        # For backwards-compatibility
+        max_attempts = 3 if max_attempts.nil? && max_duration.nil?
 
-          reload_page.refresh if reload_page
+        start_msg ||= ["with retry_until:"]
+        start_msg << "max_attempts: #{max_attempts};" if max_attempts
+        start_msg << "max_duration: #{max_duration};" if max_duration
+        start_msg << "reload_page: #{reload_page}; sleep_interval: #{sleep_interval}; raise_on_failure: #{raise_on_failure}; retry_on_exception: #{retry_on_exception}"
+        QA::Runtime::Logger.debug(start_msg.join(' '))
 
-          attempts += 1
-        end
-
-        if exit_on_failure
-          QA::Runtime::Logger.debug("Raising exception after #{max_attempts} attempts")
-          raise
+        result = nil
+        repeat_until(
+          max_attempts: max_attempts,
+          max_duration: max_duration,
+          reload_page: reload_page,
+          sleep_interval: sleep_interval,
+          raise_on_failure: raise_on_failure,
+          retry_on_exception: retry_on_exception
+        ) do
+          result = yield
         end
+        QA::Runtime::Logger.debug("ended retry_until")
 
-        false
+        result
       end
     end
   end