fix(datastore-v1): retry on subscription connection error (#2581)

lawmicha · web-flow · commit dccac9103198 · 2022-11-17T09:30:28.000-05:00
diff --git a/AmplifyPlugins/API/AWSAPICategoryPlugin/Operation/AWSGraphQLSubscriptionOperation.swift b/AmplifyPlugins/API/AWSAPICategoryPlugin/Operation/AWSGraphQLSubscriptionOperation.swift
@@ -189,6 +189,12 @@ final public class AWSGraphQLSubscriptionOperation<R: Decodable>: GraphQLSubscri
             return
         } else if case ConnectionProviderError.unauthorized = error {
             errorDescription += ": \(APIError.UnauthorizedMessageString)"
+        } else if case ConnectionProviderError.connection = error {
+            errorDescription += ": connection"
+            let error = URLError(.networkConnectionLost)
+            dispatch(result: .failure(APIError.networkError(errorDescription, nil, error)))
+            finish()
+            return
         }
 
         dispatch(result: .failure(APIError.operationError(errorDescription, "", error)))
diff --git a/AmplifyPlugins/API/AWSAPICategoryPluginTests/AWSAPICategoryPlugin+ConfigureTests.swift b/AmplifyPlugins/API/AWSAPICategoryPluginTests/AWSAPICategoryPlugin+ConfigureTests.swift
@@ -10,9 +10,10 @@ import XCTest
 @testable import Amplify
 @testable import AWSAPICategoryPlugin
 
-class AWSAPICategoryPluginConfigureTests: AWSAPICategoryPluginTestBase {
+class AWSAPICategoryPluginConfigureTests {
 
     func testPluginKey() {
+        let apiPlugin = AWSAPIPlugin()
         XCTAssertEqual(apiPlugin.key, "awsAPIPlugin")
     }
 
diff --git a/AmplifyPlugins/API/AWSAPICategoryPluginTests/AWSAPICategoryPluginTestBase.swift b/AmplifyPlugins/API/AWSAPICategoryPluginTests/AWSAPICategoryPluginTestBase.swift
@@ -30,6 +30,7 @@ class AWSAPICategoryPluginTestBase: XCTestCase {
     let testPath = "testPath"
 
     override func setUp() {
+        Amplify.reset()
         apiPlugin = AWSAPIPlugin()
 
         let authService = MockAWSAuthService()
@@ -62,7 +63,6 @@ class AWSAPICategoryPluginTestBase: XCTestCase {
             XCTFail("Failed to create endpoint config")
         }
 
-        Amplify.reset()
         let config = AmplifyConfiguration()
         do {
             try Amplify.configure(config)
@@ -72,9 +72,6 @@ class AWSAPICategoryPluginTestBase: XCTestCase {
     }
 
     override func tearDown() {
-        if let api = apiPlugin {
-            api.reset {
-            }
-        }
+        Amplify.reset()
     }
 }
diff --git a/AmplifyPlugins/API/AWSAPICategoryPluginTests/Operation/GraphQLSubscribeTests.swift b/AmplifyPlugins/API/AWSAPICategoryPluginTests/Operation/GraphQLSubscribeTests.swift
@@ -33,6 +33,8 @@ class GraphQLSubscribeTests: OperationTestBase {
     var subscriptionItem: SubscriptionItem!
     var subscriptionEventHandler: SubscriptionEventHandler!
 
+    var expectedCompletionFailureError: APIError?
+
     override func setUpWithError() throws {
         try super.setUpWithError()
 
@@ -122,7 +124,23 @@ class GraphQLSubscribeTests: OperationTestBase {
     /// - The value handler is not invoked with with a data value
     /// - The value handler is invoked with a disconnection message
     /// - The completion handler is invoked with an error termination
-    func testConnectionError() throws {
+    func testConnectionErrorWithLimitExceeded() throws {
+        receivedCompletionFinish.shouldTrigger = false
+        receivedCompletionFailure.shouldTrigger = true
+        receivedConnected.shouldTrigger = false
+        receivedDisconnected.shouldTrigger = false
+        receivedSubscriptionEventData.shouldTrigger = false
+        receivedSubscriptionEventError.shouldTrigger = false
+
+        subscribe()
+        wait(for: [onSubscribeInvoked], timeout: 0.05)
+
+        subscriptionEventHandler(.failed(ConnectionProviderError.limitExceeded(nil)), subscriptionItem)
+        expectedCompletionFailureError = APIError.operationError("", "", ConnectionProviderError.limitExceeded(nil))
+        waitForExpectations(timeout: 0.05)
+    }
+
+    func testConnectionErrorWithSubscriptionError() throws {
         receivedCompletionFinish.shouldTrigger = false
         receivedCompletionFailure.shouldTrigger = true
         receivedConnected.shouldTrigger = false
@@ -134,8 +152,42 @@ class GraphQLSubscribeTests: OperationTestBase {
         wait(for: [onSubscribeInvoked], timeout: 0.05)
 
         subscriptionEventHandler(.connection(.connecting), subscriptionItem)
-        subscriptionEventHandler(.failed("Error"), subscriptionItem)
+        subscriptionEventHandler(.failed(ConnectionProviderError.subscription("", nil)), subscriptionItem)
+        expectedCompletionFailureError = APIError.operationError("", "", ConnectionProviderError.subscription("", nil))
+        waitForExpectations(timeout: 0.05)
+    }
+
+    func testConnectionErrorWithConnectionUnauthorizedError() throws {
+        receivedCompletionFinish.shouldTrigger = false
+        receivedCompletionFailure.shouldTrigger = true
+        receivedConnected.shouldTrigger = false
+        receivedDisconnected.shouldTrigger = false
+        receivedSubscriptionEventData.shouldTrigger = false
+        receivedSubscriptionEventError.shouldTrigger = false
 
+        subscribe()
+        wait(for: [onSubscribeInvoked], timeout: 0.05)
+
+        subscriptionEventHandler(.connection(.connecting), subscriptionItem)
+        subscriptionEventHandler(.failed(ConnectionProviderError.unauthorized), subscriptionItem)
+        expectedCompletionFailureError = APIError.operationError("", "", ConnectionProviderError.unauthorized)
+        waitForExpectations(timeout: 0.05)
+    }
+
+    func testConnectionErrorWithConnectionProviderConnectionError() throws {
+        receivedCompletionFinish.shouldTrigger = false
+        receivedCompletionFailure.shouldTrigger = true
+        receivedConnected.shouldTrigger = false
+        receivedDisconnected.shouldTrigger = false
+        receivedSubscriptionEventData.shouldTrigger = false
+        receivedSubscriptionEventError.shouldTrigger = false
+
+        subscribe()
+        wait(for: [onSubscribeInvoked], timeout: 0.05)
+
+        subscriptionEventHandler(.connection(.connecting), subscriptionItem)
+        subscriptionEventHandler(.failed(ConnectionProviderError.connection), subscriptionItem)
+        expectedCompletionFailureError = APIError.networkError("", nil, URLError(.networkConnectionLost))
         waitForExpectations(timeout: 0.05)
     }
 
@@ -286,7 +338,11 @@ class GraphQLSubscribeTests: OperationTestBase {
                 }
         }, completionListener: { result in
             switch result {
-            case .failure:
+            case .failure(let error):
+                if let apiError = error as? APIError,
+                   let expectedError = self.expectedCompletionFailureError {
+                    XCTAssertEqual(apiError, expectedError)
+                }
                 self.receivedCompletionFailure.fulfill()
             case .success:
                 self.receivedCompletionFinish.fulfill()
@@ -296,3 +352,43 @@ class GraphQLSubscribeTests: OperationTestBase {
         return operation
     }
 }
+
+extension APIError: Equatable {
+    public static func == (lhs: APIError, rhs: APIError) -> Bool {
+        switch (lhs, rhs) {
+        case (.unknown, .unknown),
+            (.invalidConfiguration, .invalidConfiguration),
+            (.httpStatusError, .httpStatusError),
+            (.pluginError, .pluginError):
+            return true
+        case (.operationError(_, _, let lhs), .operationError(_, _, let rhs)):
+            if let lhs = lhs as? ConnectionProviderError, let rhs = rhs as? ConnectionProviderError {
+                switch (lhs, rhs) {
+                case (.connection, .connection),
+                    (.jsonParse, .jsonParse),
+                    (.limitExceeded, .limitExceeded),
+                    (.subscription, .subscription),
+                    (.unauthorized, .unauthorized),
+                    (.unknown, .unknown):
+                    return true
+                default:
+                    return false
+                }
+            } else if lhs == nil && rhs == nil {
+                return true
+            } else {
+                return false
+            }
+        case (.networkError(_, _, let lhs), .networkError(_, _, let rhs)):
+            if let lhs = lhs as? URLError, let rhs = rhs as? URLError {
+                return lhs.code == rhs.code
+            } else if lhs == nil && rhs == nil {
+                return true
+            } else {
+                return false
+            }
+        default:
+            return false
+        }
+    }
+}
diff --git a/AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Sync/RemoteSyncEngine+Retryable.swift b/AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Sync/RemoteSyncEngine+Retryable.swift
@@ -36,6 +36,12 @@ extension RemoteSyncEngine {
             urlErrorOptional = underlyingError
         } else if let urlError = error as? URLError {
             urlErrorOptional = urlError
+        } else if let dataStoreError = error as? DataStoreError,
+                  case .api(let amplifyError, _) = dataStoreError,
+                  let apiError = amplifyError as? APIError,
+                  case .networkError(_, _, let error) = apiError,
+                  let urlError = error as? URLError {
+            urlErrorOptional = urlError
         }
 
         let advice = requestRetryablePolicy.retryRequestAdvice(urlError: urlErrorOptional,
diff --git a/AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Sync/RequestRetryablePolicy.swift b/AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Sync/RequestRetryablePolicy.swift
@@ -39,7 +39,8 @@ class RequestRetryablePolicy: RequestRetryable {
              .cannotFindHost,
              .timedOut,
              .dataNotAllowed,
-             .cannotParseResponse:
+             .cannotParseResponse,
+             .networkConnectionLost:
             let waitMillis = retryDelayInMillseconds(for: attemptNumber)
             return RequestRetryAdvice(shouldRetry: true, retryInterval: .milliseconds(waitMillis))
         default:
diff --git a/AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginTests/Sync/RequestRetryablePolicyTests.swift b/AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginTests/Sync/RequestRetryablePolicyTests.swift
@@ -188,6 +188,17 @@ class RequestRetryablePolicyTests: XCTestCase {
         assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
     }
 
+    func testNetworkConnectionLostError() {
+        let retryableErrorCode = URLError.init(.networkConnectionLost)
+
+        let retryAdvice = retryPolicy.retryRequestAdvice(urlError: retryableErrorCode,
+                                                         httpURLResponse: nil,
+                                                         attemptNumber: 1)
+
+        XCTAssert(retryAdvice.shouldRetry)
+        assertMilliseconds(retryAdvice.retryInterval, greaterThan: 200, lessThan: 300)
+    }
+
     func testHTTPTooManyRedirectsError() {
         let nonRetryableErrorCode = URLError.init(.httpTooManyRedirects)
 

Original file line number	Diff line number	Diff line change
`@@ -10,9 +10,10 @@ import XCTest`
`10`	`10`	`@testable import Amplify`
`11`	`11`	`@testable import AWSAPICategoryPlugin`
`12`	`12`
`13`		`-class AWSAPICategoryPluginConfigureTests: AWSAPICategoryPluginTestBase {`
	`13`	`+class AWSAPICategoryPluginConfigureTests {`
`14`	`14`
`15`	`15`	`func testPluginKey() {`
	`16`	`+ let apiPlugin = AWSAPIPlugin()`
`16`	`17`	`XCTAssertEqual(apiPlugin.key, "awsAPIPlugin")`
`17`	`18`	`}`
`18`	`19`
Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,7 @@ class AWSAPICategoryPluginTestBase: XCTestCase {`
`30`	`30`	`let testPath = "testPath"`
`31`	`31`
`32`	`32`	`override func setUp() {`
	`33`	`+ Amplify.reset()`
`33`	`34`	`apiPlugin = AWSAPIPlugin()`
`34`	`35`
`35`	`36`	`let authService = MockAWSAuthService()`
`@@ -62,7 +63,6 @@ class AWSAPICategoryPluginTestBase: XCTestCase {`
`62`	`63`	`XCTFail("Failed to create endpoint config")`
`63`	`64`	`}`
`64`	`65`
`65`		`- Amplify.reset()`
`66`	`66`	`let config = AmplifyConfiguration()`
`67`	`67`	`do {`
`68`	`68`	`try Amplify.configure(config)`
`@@ -72,9 +72,6 @@ class AWSAPICategoryPluginTestBase: XCTestCase {`
`72`	`72`	`}`
`73`	`73`
`74`	`74`	`override func tearDown() {`
`75`		`- if let api = apiPlugin {`
`76`		`- api.reset {`
`77`		`- }`
`78`		`- }`
	`75`	`+ Amplify.reset()`
`79`	`76`	`}`
`80`	`77`	`}`