Ensure push notifications re-register after store switch

WooCommerce/Classes/System/ProcessConfiguration.swift

@@ -42,3 +42,10 @@ struct ProcessConfiguration {
         ProcessInfo.processInfo.arguments.contains("use-mocked-card-present-payment")
     }
 }
+
+extension ProcessInfo {
+    /// Indicates whether the current process is executing under XCTest.
+    static var isRunningUnitTests: Bool {
+        NSClassFromString("XCTestCase") != nil
+    }
+}

WooCommerce/Classes/Yosemite/DefaultStoresManager.swift

@@ -357,6 +357,8 @@ class DefaultStoresManager: StoresManager {
         restoreSessionSiteIfPossible()
         ServiceLocator.pushNotesManager.reloadBadgeCount()
 
+        registerForPushNotificationsIfNeeded()
+
         NotificationCenter.default.post(name: .StoresManagerDidUpdateDefaultSite, object: nil)
     }
 
@@ -402,6 +404,24 @@ class DefaultStoresManager: StoresManager {
 //
 private extension DefaultStoresManager {
 
+    func registerForPushNotificationsIfNeeded() {
+        #if targetEnvironment(simulator)
+        let shouldRegisterForPushNotifications = ProcessInfo.isRunningUnitTests
+        #else
+        let shouldRegisterForPushNotifications = true
+        #endif
+
+        guard shouldRegisterForPushNotifications else {
+            return
+        }
+
+        DispatchQueue.main.async {
+            let pushNotesManager = ServiceLocator.pushNotesManager
+            pushNotesManager.registerForRemoteNotifications()
+            pushNotesManager.ensureAuthorizationIsRequested(includesProvisionalAuth: false, onCompletion: nil)
+        }
+    }
+
     /// Loads the Default Account into the current Session, if possible.
     ///
     func restoreSessionAccountIfPossible() {

WooCommerce/WooCommerceTests/Mocks/MockPushNotificationsManager.swift

@@ -57,6 +57,10 @@ final class MockPushNotificationsManager: PushNotesManager {
     private(set) var canceledLocalNotificationScenarios: [[LocalNotification.Scenario]] = []
     private(set) var resetBadgeCountKinds: [Note.Kind] = []
     var onRequestLocalNotificationCalled: (() -> Void)?
+    private(set) var registerForRemoteNotificationsCallCount = 0
+    var onRegisterForRemoteNotifications: (() -> Void)?
+    private(set) var ensureAuthorizationIsRequestedCallCount = 0
+    var onEnsureAuthorizationIsRequested: (() -> Void)?
 
     init(mockedDeviceID: String? = nil) {
         self.mockedDeviceID = mockedDeviceID
@@ -75,15 +79,18 @@ final class MockPushNotificationsManager: PushNotesManager {
     }
 
     func registerForRemoteNotifications() {
-
+        registerForRemoteNotificationsCallCount += 1
+        onRegisterForRemoteNotifications?()
     }
 
     func unregisterForRemoteNotifications() {
 
     }
 
     func ensureAuthorizationIsRequested(includesProvisionalAuth: Bool, onCompletion: ((Bool) -> ())?) {
-
+        ensureAuthorizationIsRequestedCallCount += 1
+        onEnsureAuthorizationIsRequested?()
+        onCompletion?(true)
     }
 
     func registrationDidFail(with error: Error) {

WooCommerce/WooCommerceTests/Yosemite/StoresManagerTests.swift

@@ -246,6 +246,28 @@ final class StoresManagerTests: XCTestCase {
         XCTAssertEqual(siteIDValues, [nil, siteID, nil])
     }
 
+    func test_updateDefaultStore_registersForRemoteNotifications() {
+        // Arrange
+        let pushNotificationsManager = MockPushNotificationsManager()
+        ServiceLocator.setPushNotesManager(pushNotificationsManager)
+        defer {
+            ServiceLocator.setPushNotesManager(PushNotificationsManager())
+        }
+        let manager = DefaultStoresManager.testingInstance
+        let expectation = expectation(description: "register called")
+        pushNotificationsManager.onRegisterForRemoteNotifications = {
+            expectation.fulfill()
+        }
+
+        // Action
+        manager.updateDefaultStore(storeID: 1_234)
+
+        // Assert
+        wait(for: [expectation], timeout: 1.0)
+        XCTAssertEqual(pushNotificationsManager.registerForRemoteNotificationsCallCount, 1)
+        XCTAssertEqual(pushNotificationsManager.ensureAuthorizationIsRequestedCallCount, 1)
+    }
+
     // MARK: `updateDefaultStore(_ site: Site)`
 
     func test_updateDefaultStore_with_the_same_siteID_updates_site_but_does_not_emit_siteID() {