fix: Add retry mechanism to flakey tests due to IAM eventual consistency (#2061)

Author: dayaffe

IntegrationTests/Services/AWSCognitoIdentityIntegrationTests/CognitoAWSCredentialIdentityResolverTests.swift

@@ -95,16 +95,31 @@ class CognitoAWSCredentialIdentityResolverTests: XCTestCase {
             region: "us-east-1" // different from cognito pool region
         )
         let cognitoStsClient = STSClient(config: cognitoStsConfig)
-        let response = try await cognitoStsClient.getCallerIdentity(
-            input: GetCallerIdentityInput()
-        )
 
-        let account = try XCTUnwrap(response.account)
-        let userId = try XCTUnwrap(response.userId)
-        let arn = try XCTUnwrap(response.arn)
+        // Retry to handle IAM eventual consistency
+        var lastError: Error?
+        let totalRetries = 5
+        for attempt in 0..<totalRetries {
+            do {
+                let response = try await cognitoStsClient.getCallerIdentity(
+                    input: GetCallerIdentityInput()
+                )
+
+                let account = try XCTUnwrap(response.account)
+                let userId = try XCTUnwrap(response.userId)
+                let arn = try XCTUnwrap(response.arn)
 
-        XCTAssertNotEqual(account, "")
-        XCTAssertNotEqual(userId, "")
-        XCTAssertTrue(arn.contains(roleName))
+                XCTAssertNotEqual(account, "")
+                XCTAssertNotEqual(userId, "")
+                XCTAssertTrue(arn.contains(roleName))
+                return
+            } catch {
+                lastError = error
+                if attempt < (totalRetries-1) {
+                    try await Task.sleep(nanoseconds: 3_000_000_000) // 3 seconds
+                }
+            }
+        }
+        throw lastError!
     }
 }

IntegrationTests/Services/AWSSTSIntegrationTests/STSAssumeRoleAWSCredentialIdentityResolverTests.swift

@@ -62,19 +62,31 @@ class STSAssumeRoleAWSCredentialIdentityResolverTests: XCTestCase {
 
     // Confirm STS assume role credentials provider works by validating response.
     func testGetCallerIdentity() async throws {
-        let response = try await assumeRoleStsClient.getCallerIdentity(
-            input: GetCallerIdentityInput()
-        )
-
-        // Ensure returned caller info aren't nil
-        let account = try XCTUnwrap(response.account)
-        let userId = try XCTUnwrap(response.userId)
-        let arn = try XCTUnwrap(response.arn)
-
-        // Ensure returned caller info aren't empty strings
-        XCTAssertNotEqual(account, "")
-        XCTAssertNotEqual(userId, "")
-        XCTAssertNotEqual(arn, "")
+        // Retry to handle IAM eventual consistency
+        var lastError: Error?
+        let totalRetries = 5
+        for attempt in 0..<totalRetries {
+            do {
+                let response = try await assumeRoleStsClient.getCallerIdentity(
+                    input: GetCallerIdentityInput()
+                )
+
+                let account = try XCTUnwrap(response.account)
+                let userId = try XCTUnwrap(response.userId)
+                let arn = try XCTUnwrap(response.arn)
+
+                XCTAssertNotEqual(account, "")
+                XCTAssertNotEqual(userId, "")
+                XCTAssertNotEqual(arn, "")
+                return
+            } catch {
+                lastError = error
+                if attempt < (totalRetries-1) {
+                    try await Task.sleep(nanoseconds: 3_000_000_000) // 3 seconds
+                }
+            }
+        }
+        throw lastError!
     }
 
     // Right now opentelemetry-swift doesnt support linux or visionos

IntegrationTests/Services/AWSSTSIntegrationTests/STSWebIdentityAWSCredentialIdentityResolverTests.swift

@@ -136,19 +136,31 @@ class STSWebIdentityAWSCredentialIdentityResolverTests: XCTestCase {
 
     // Confirm STS web identity credentials provider works by validating response.
     func testGetCallerIdentity() async throws {
-        let response = try await webIdentityStsClient.getCallerIdentity(
-            input: GetCallerIdentityInput()
-        )
-
-        // Ensure returned caller info aren't nil
-        let account = try XCTUnwrap(response.account)
-        let userId = try XCTUnwrap(response.userId)
-        let arn = try XCTUnwrap(response.arn)
-
-        // Ensure returned caller info aren't empty strings
-        XCTAssertNotEqual(account, "")
-        XCTAssertNotEqual(userId, "")
-        XCTAssertNotEqual(arn, "")
+        // Retry to handle IAM eventual consistency
+        var lastError: Error?
+        let totalRetries = 5
+        for attempt in 0..<totalRetries {
+            do {
+                let response = try await webIdentityStsClient.getCallerIdentity(
+                    input: GetCallerIdentityInput()
+                )
+
+                let account = try XCTUnwrap(response.account)
+                let userId = try XCTUnwrap(response.userId)
+                let arn = try XCTUnwrap(response.arn)
+
+                XCTAssertNotEqual(account, "")
+                XCTAssertNotEqual(userId, "")
+                XCTAssertNotEqual(arn, "")
+                return
+            } catch {
+                lastError = error
+                if attempt < (totalRetries-1) {
+                    try await Task.sleep(nanoseconds: 3_000_000_000) // 3 seconds
+                }
+            }
+        }
+        throw lastError!
     }
 
     // MARK: - SETUP HELPER FUNCTIONS