fix(DataStore): endless retry of mutation request when server responds with 401 error code (#3511) (#3512)

* fix(DataStore): endless retry of mutation request when server responds with 401 error code (#3511)

* fix(DataStore): endless retry of mutation request when server responds with 401 error code (#3511)

* fix(DataStore): endless retry of mutation request when server  responds with 401 error code (#3511)

Co-authored-by: Michael Law <[email protected]>

---------

Co-authored-by: Michael Law <[email protected]>

Author: MuniekMg

AmplifyPlugins/Core/AWSPluginsCore/Auth/AWSAuthModeStrategy.swift

@@ -55,6 +55,9 @@ public protocol AuthorizationTypeIterator {
 
     /// Total number of values
     var count: Int { get }
+    
+    /// Whether iterator has next available `AuthorizationType` to return or not
+    var hasNext: Bool { get }
 
     /// Next available `AuthorizationType` or `nil` if exhausted
     mutating func next() -> AuthorizationType?
@@ -66,18 +69,29 @@ public struct AWSAuthorizationTypeIterator: AuthorizationTypeIterator {
 
     private var values: IndexingIterator<[AWSAuthorizationType]>
     private var _count: Int
+    private var _position: Int
 
     public init(withValues values: [AWSAuthorizationType]) {
         self.values = values.makeIterator()
         self._count = values.count
+        self._position = 0
     }
 
     public var count: Int {
         _count
     }
+    
+    public var hasNext: Bool {
+        _position < _count
+    }
 
     public mutating func next() -> AWSAuthorizationType? {
-        values.next()
+        if let value = values.next() {
+            _position += 1
+            return value
+        }
+        
+        return nil
     }
 }
 

AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Sync/MutationSync/OutgoingMutationQueue/SyncMutationToCloudOperation.swift

@@ -321,7 +321,7 @@ class SyncMutationToCloudOperation: AsynchronousOperation {
 
     /// - Warning: Must be invoked from a locking context
     private func shouldRetryWithDifferentAuthType() -> RequestRetryAdvice {
-        let shouldRetry = (authTypesIterator?.count ?? 0) > 0
+        let shouldRetry = authTypesIterator?.hasNext == true
         return RequestRetryAdvice(shouldRetry: shouldRetry, retryInterval: .milliseconds(0))
     }
 

AmplifyPlugins/DataStore/Tests/AWSDataStorePluginTests/Sync/MutationQueue/SyncMutationToCloudOperationTests.swift

@@ -292,6 +292,68 @@ class SyncMutationToCloudOperationTests: XCTestCase {
         XCTAssertTrue(advice.shouldRetry)
     }
 
+    /// Given: Model with multiple auth types. Mutation requests always fail with 401 error code
+    /// When: Mutating model fails with 401
+    /// Then: DataStore will try again with each auth type and eventually fails
+    func testGetRetryAdviceForEachModelAuthTypeThenFail_HTTPStatusError401() async throws {
+        var numberOfTimesEntered = 0
+        let mutationEvent = try createMutationEvent()
+        let authStrategy = MockMultiAuthModeStrategy()
+        let expectedNumberOfTimesEntered = authStrategy.authTypesFor(schema: mutationEvent.schema, operation: .create).count
+        
+        let expectCalllToApiMutateNTimesAndFail = expectation(description: "Call API.mutate \(expectedNumberOfTimesEntered) times and then fail")
+        
+        let response = HTTPURLResponse(url: URL(string: "http://localhost")!,
+                                       statusCode: 401,
+                                       httpVersion: nil,
+                                       headerFields: nil)!
+        let error = APIError.httpStatusError(401, response)
+        
+        let operation = await SyncMutationToCloudOperation(
+            mutationEvent: mutationEvent,
+            getLatestSyncMetadata: { nil },
+            api: mockAPIPlugin,
+            authModeStrategy: authStrategy,
+            networkReachabilityPublisher: publisher,
+            currentAttemptNumber: 1,
+            completion: { result in
+                if numberOfTimesEntered == expectedNumberOfTimesEntered {
+                    expectCalllToApiMutateNTimesAndFail.fulfill()
+                    
+                } else {
+                    XCTFail("API.mutate was called incorrect amount of times, expected: \(expectedNumberOfTimesEntered), was : \(numberOfTimesEntered)")
+                }
+            }
+        )
+        
+        let responder = MutateRequestListenerResponder<MutationSync<AnyModel>> { request, eventListener in
+            let requestOptions = GraphQLOperationRequest<MutationSync<AnyModel>>.Options(pluginOptions: nil)
+            let request = GraphQLOperationRequest<MutationSync<AnyModel>>(apiName: request.apiName,
+                                                     operationType: .mutation,
+                                                     document: request.document,
+                                                     variables: request.variables,
+                                                     responseType: request.responseType,
+                                                     options: requestOptions)
+            let operation = MockGraphQLOperation(request: request, responseType: request.responseType)
+            
+            numberOfTimesEntered += 1
+            
+            DispatchQueue.global().sync {
+                // Fail with 401 status code
+                eventListener!(.failure(error))
+            }
+            
+            return operation
+        }
+        
+        mockAPIPlugin.responders[.mutateRequestListener] = responder
+        
+        let queue = OperationQueue()
+        queue.addOperation(operation)
+        
+        await fulfillment(of: [expectCalllToApiMutateNTimesAndFail], timeout: defaultAsyncWaitTimeout)
+    }
+    
     func testGetRetryAdvice_OperationErrorAuthErrorWithMultiAuth_RetryTrue() async throws {
         let operation = await SyncMutationToCloudOperation(
             mutationEvent: try createMutationEvent(),

AmplifyPlugins/DataStore/Tests/AWSDataStorePluginTests/Sync/Support/AWSAuthorizationTypeIteratorTests.swift

@@ -0,0 +1,40 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import XCTest
+import AWSPluginsCore
+
+class AWSAuthorizationTypeIteratorTests: XCTestCase {
+    
+    func testEmptyIterator_hasNextValue_false() throws {
+        var iterator = AWSAuthorizationTypeIterator(withValues: [])
+       
+        XCTAssertFalse(iterator.hasNext)
+        XCTAssertNil(iterator.next())
+    }
+    
+    func testOneElementIterator_hasNextValue_once() throws {
+        var iterator = AWSAuthorizationTypeIterator(withValues: [.amazonCognitoUserPools])
+        
+        XCTAssertTrue(iterator.hasNext)
+        XCTAssertNotNil(iterator.next())
+        
+        XCTAssertFalse(iterator.hasNext)
+    }
+    
+    func testTwoElementsIterator_hasNextValue_twice() throws {
+        var iterator = AWSAuthorizationTypeIterator(withValues: [.amazonCognitoUserPools, .apiKey])
+        
+        XCTAssertTrue(iterator.hasNext)
+        XCTAssertNotNil(iterator.next())
+        
+        XCTAssertTrue(iterator.hasNext)
+        XCTAssertNotNil(iterator.next())
+        
+        XCTAssertFalse(iterator.hasNext)
+    }
+}