Add extra tests for retry behavior

gems/aws-sdk-core/spec/aws/plugins/retry_errors_spec.rb

@@ -272,6 +272,62 @@ module Plugins
           handle_with_retry(test_case_def)
         end
 
+        it 'fails due to retry quota bucket exhaustion' do
+          config.max_attempts = 5
+          config.retry_quota.instance_variable_set(:@available_capacity, 10)
+
+          test_case_def = [
+            {
+              response: { status_code: 500, error: service_error },
+              expect: { available_capacity: 5, retries: 1, delay: 1 }
+            },
+            {
+              response: { status_code: 502, error: service_error },
+              expect: { available_capacity: 0, retries: 2, delay: 2 }
+            },
+            {
+              response: { status_code: 503, error: service_error },
+              expect: { available_capacity: 0, retries: 2 }
+            }
+          ]
+
+          handle_with_retry(test_case_def)
+        end
+
+        it 'recovers after successful responses' do
+          config.max_attempts = 5
+          config.retry_quota.instance_variable_set(:@available_capacity, 15)
+
+          test_case_def = [
+            {
+              response: { status_code: 500, error: service_error },
+              expect: { available_capacity: 10, retries: 1, delay: 1 }
+            },
+            {
+              response: { status_code: 502, error: service_error },
+              expect: { available_capacity: 5, retries: 2, delay: 2 }
+            },
+            {
+              response: { status_code: 200, error: nil },
+              expect: { available_capacity: 10, retries: 2 }
+            }
+          ]
+          handle_with_retry(test_case_def)
+
+          test_case_post_success = [
+            {
+              response: { status_code: 500, error: service_error },
+              expect: { available_capacity: 5, retries: 1, delay: 1 }
+            },
+            {
+              response: { status_code: 200, error: nil },
+              expect: { available_capacity: 10, retries: 1 }
+            }
+          ]
+          reset_request
+          handle_with_retry(test_case_post_success)
+        end
+
         it 'corrects and retries clock skew errors' do
           clock_skew_error = RetryErrorsSvc::Errors::RequestTimeTooSkewed
                                .new(nil, nil)

gems/aws-sdk-core/spec/retry_errors_helper.rb

@@ -17,16 +17,10 @@ def handle(send_handler = nil, &block)
 # expect: delay, available_capacity, retries, calculated_rate,
 #   measured_tx_rate, fill_rate
 def handle_with_retry(test_cases)
-  # Apply delay expectations first
-  test_cases.each do |test_case|
-    if test_case[:expect][:delay]
-      expect(Kernel).to receive(:sleep).with(test_case[:expect][:delay])
-    end
-  end
-
   i = 0
   handle do |_context|
-    apply_expectations(test_cases[i - 1]) if i > 0
+    apply_delay(test_cases[i])
+    apply_expectations(test_cases[i - 1]) if i.positive?
 
     # Setup the next response
     setup_next_response(test_cases[i])
@@ -38,15 +32,30 @@ def handle_with_retry(test_cases)
 
   expect(i).to(
     eq(test_cases.size),
-    "Wrong number of retries. Handler was called #{i} times but "\
+    "Wrong number of retries. Handler was called #{i} times but " \
     "#{test_cases.size} test cases were defined."
   )
 
   # Handle has finished called.  Apply final expectations.
   apply_expectations(test_cases[i - 1])
 end
 
-# apply the expectations from a test case
+# Reset the request context for a subsequent call
+def reset_request
+  resp.context.retries = 0
+  resp.context.metadata[:retries] = {}
+end
+
+# apply a delay to the current test case
+# See handle_with_retry for test case definition
+def apply_delay(test_case)
+  expected = test_case[:expect]
+  return unless expected[:delay]
+
+  expect(Kernel).to receive(:sleep).with(expected[:delay])
+end
+
+# apply the expectations from a previous test case
 # See handle_with_retry for test case definition
 def apply_expectations(test_case)
   expected = test_case[:expect]
@@ -95,6 +104,8 @@ def apply_expectations(test_case)
   end
 end
 
+# setup the next response for the handler
+# See handle_with_retry for test case definition
 def setup_next_response(test_case)
   response = test_case[:response]
   resp.context.http_response.status_code = response[:status_code]