From 565fff74b58913c17e0ebfc3bbc8479cf4df84f0 Mon Sep 17 00:00:00 2001
From: Matt Muller <mamuller@amazon.com>
Date: Mon, 28 Oct 2024 18:27:00 -0400
Subject: [PATCH] Add extra tests for retry behavior

---
 .../spec/aws/plugins/retry_errors_spec.rb     | 47 +++++++++++++++++++
 gems/aws-sdk-core/spec/retry_errors_helper.rb | 29 ++++++------
 2 files changed, 61 insertions(+), 15 deletions(-)

diff --git a/gems/aws-sdk-core/spec/aws/plugins/retry_errors_spec.rb b/gems/aws-sdk-core/spec/aws/plugins/retry_errors_spec.rb
index 9b56972f352..277e7dfcb2f 100644
--- a/gems/aws-sdk-core/spec/aws/plugins/retry_errors_spec.rb
+++ b/gems/aws-sdk-core/spec/aws/plugins/retry_errors_spec.rb
@@ -272,6 +272,53 @@ 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 }
+            },
+            {
+              response: { status_code: 500, error: service_error },
+              expect: { available_capacity: 5, retries: 1, delay: 1 }
+            }
+          ]
+          handle_with_retry(test_case_def)
+        end
+
         it 'corrects and retries clock skew errors' do
           clock_skew_error = RetryErrorsSvc::Errors::RequestTimeTooSkewed
                                .new(nil, nil)
diff --git a/gems/aws-sdk-core/spec/retry_errors_helper.rb b/gems/aws-sdk-core/spec/retry_errors_helper.rb
index 93497664c05..759a675ba88 100644
--- a/gems/aws-sdk-core/spec/retry_errors_helper.rb
+++ b/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])
@@ -36,17 +30,20 @@ def handle_with_retry(test_cases)
     resp
   end
 
-  expect(i).to(
-    eq(test_cases.size),
-    "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
+# 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 +92,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]