Ensure that unregistering push notifications completes

Author: itsmeichigo

WooCommerce/Classes/Notifications/PushNotificationsManager.swift

@@ -167,18 +167,19 @@ extension PushNotificationsManager {
 
     /// Unregisters the Application from WordPress.com Push Notifications Service.
     ///
-    func unregisterForRemoteNotifications() {
+    func unregisterForRemoteNotifications(onCompletion: @escaping () -> Void) {
         DDLogInfo("📱 Unregistering For Remote Notifications...")
 
         unregisterDotcomDeviceIfPossible() { error in
             if let error = error {
                 DDLogError("⛔️ Unable to unregister from WordPress.com Push Notifications: \(error)")
-                return
+            } else {
+                DDLogInfo("📱 Successfully unregistered from WordPress.com Push Notifications!")
+                self.deviceID = nil
+                self.deviceToken = nil
             }
-
-            DDLogInfo("📱 Successfully unregistered from WordPress.com Push Notifications!")
-            self.deviceID = nil
-            self.deviceToken = nil
+            // Always call completion, even on error
+            onCompletion()
         }
     }
 
@@ -249,7 +250,7 @@ extension PushNotificationsManager {
     ///
     func registrationDidFail(with error: Error) {
         DDLogError("⛔️ Push Notifications Registration Failure: \(error)")
-        unregisterForRemoteNotifications()
+        unregisterForRemoteNotifications {}
     }
 
     /// Handles a notification while the app is in foreground

WooCommerce/Classes/ServiceLocator/PushNotesManager.swift

@@ -51,7 +51,7 @@ protocol PushNotesManager {
 
     /// Unregisters the Application from WordPress.com Push Notifications Service.
     ///
-    func unregisterForRemoteNotifications()
+    func unregisterForRemoteNotifications(onCompletion: @escaping () -> Void)
 
     /// Requests Authorization to receive Push Notifications, *only* when the current Status is not determined.
     ///

WooCommerce/Classes/Yosemite/AuthenticatedState.swift

@@ -237,11 +237,6 @@ class AuthenticatedState: StoresManagerState {
     /// Executed before the current state is deactivated.
     ///
     func willLeave() {
-        let pushNotesManager = ServiceLocator.pushNotesManager
-
-        pushNotesManager.unregisterForRemoteNotifications()
-        pushNotesManager.resetBadgeCountForAllStores(onCompletion: {})
-
         resetServices()
     }
 

WooCommerce/Classes/Yosemite/DefaultStoresManager.swift

@@ -270,7 +270,7 @@ class DefaultStoresManager: StoresManager {
         sessionManager.deleteApplicationPassword(locally: true)
         ServiceLocator.analytics.refreshUserData()
         ZendeskProvider.shared.reset()
-        ServiceLocator.pushNotesManager.unregisterForRemoteNotifications()
+        ServiceLocator.pushNotesManager.unregisterForRemoteNotifications {}
     }
 
     /// Fully deauthenticates the user, if needed.
@@ -290,6 +290,19 @@ class DefaultStoresManager: StoresManager {
     ///
     @discardableResult
     func deauthenticate() -> StoresManager {
+        let pushNotesManager = ServiceLocator.pushNotesManager
+        pushNotesManager.resetBadgeCountForAllStores(onCompletion: {})
+
+        // Keep a strong reference to the current state to prevent it from being deallocated
+        // until the unregister completes
+        let currentState = state
+
+        // Unregister from remote notifications asynchronously
+        pushNotesManager.unregisterForRemoteNotifications {
+            // Release the strong reference to allow state cleanup
+            _ = currentState
+        }
+
         applicationPasswordGenerationFailureObserver = nil
         invalidWPCOMTokenNotificationObserver = nil
         stopObservingNetworkNotifications()

WooCommerce/WooCommerceTests/Notifications/PushNotificationsManagerTests.swift

@@ -134,7 +134,7 @@ final class PushNotificationsManagerTests: XCTestCase {
     ///
     func testUnregisterForRemoteNotificationsDoesNothingWhenThereIsNoDeviceIdStored() {
         XCTAssert(storesManager.receivedActions.isEmpty)
-        manager.unregisterForRemoteNotifications()
+        manager.unregisterForRemoteNotifications {}
         XCTAssert(storesManager.receivedActions.isEmpty)
     }
 
@@ -143,7 +143,7 @@ final class PushNotificationsManagerTests: XCTestCase {
     ///
     func testUnregisterForRemoteNotificationsEffectivelyDispatchesUnregisterDeviceAction() {
         defaults.set(Sample.deviceID, forKey: .deviceID)
-        manager.unregisterForRemoteNotifications()
+        manager.unregisterForRemoteNotifications {}
 
         guard case let .unregisterDevice(deviceID, _) = storesManager.receivedActions.first as! NotificationAction else {
             XCTFail()
@@ -161,7 +161,7 @@ final class PushNotificationsManagerTests: XCTestCase {
         defaults.set(Sample.deviceID, forKey: .deviceID)
         defaults.set(Sample.deviceToken, forKey: .deviceToken)
 
-        manager.unregisterForRemoteNotifications()
+        manager.unregisterForRemoteNotifications {}
 
         guard case let .unregisterDevice(_, onCompletion) = storesManager.receivedActions.first as! NotificationAction else {
             XCTFail()