Watch for badge count changes and update tab bar badge count accordingly

WordPress/Classes/NotificationsManager.m

@@ -87,18 +87,20 @@ + (BOOL)deviceRegisteredForPushNotifications {
 + (void)handleNotification:(NSDictionary *)userInfo forState:(UIApplicationState)state completionHandler:(void (^)(UIBackgroundFetchResult))completionHandler {
     DDLogVerbose(@"Received push notification:\nPayload: %@\nCurrent Application state: %d", userInfo, state);
 
+    // Try to pull the badge number from the notification object
+    // Badge count does not normally update when the app is active
+    // And this forces KVO to be fired
+    NSDictionary *apsObject = [userInfo dictionaryForKey:@"aps"];
+    if (apsObject) {
+        NSNumber *badgeCount = [apsObject numberForKey:@"badge"];
+        if (badgeCount) {
+            [UIApplication sharedApplication].applicationIconBadgeNumber = [badgeCount intValue];
+        }
+    }
+
     if ([userInfo stringForKey:@"type"]) { //check if it is the badge reset PN
         NSString *notificationType = [userInfo stringForKey:@"type"];
         if ([notificationType isEqualToString:@"badge-reset"]) {
-            [UIApplication sharedApplication].applicationIconBadgeNumber = 0;
-            //Try to pull the badge number from the notification object
-            NSDictionary *apsObject = [userInfo dictionaryForKey:@"aps"];
-            if (apsObject) {
-                NSNumber *badgeCount = [apsObject numberForKey:@"badge"];
-                if (badgeCount) {
-                    [UIApplication sharedApplication].applicationIconBadgeNumber = [badgeCount intValue];
-                }
-            }
             return;
         }
     }

WordPress/Classes/NotificationsViewController.m

@@ -92,6 +92,7 @@ - (BOOL)showJetpackConnectMessage {
 
 - (void)dealloc {
     [[NSNotificationCenter defaultCenter] removeObserver:self];
+    [[UIApplication sharedApplication] removeObserver:self forKeyPath:@"applicationIconBadgeNumber"];
 }
 
 - (void)viewDidLoad
@@ -111,6 +112,14 @@ - (void)viewDidLoad
                                                                         action:@selector(showNotificationSettings)];
         self.navigationItem.rightBarButtonItem = pushSettings;
     }
+
+    // Watch for application badge number changes
+    UIApplication *application = [UIApplication sharedApplication];
+    [application addObserver:self
+                  forKeyPath:@"applicationIconBadgeNumber"
+                     options:NSKeyValueObservingOptionNew
+                     context:nil];
+    [self updateTabBarBadgeNumber];
 }
 
 - (void)viewWillAppear:(BOOL)animated {
@@ -140,9 +149,27 @@ - (void)viewDidDisappear:(BOOL)animated {
         [self pruneOldNotes];
 }
 
+#pragma mark - NSObject(NSKeyValueObserving) methods
+
+- (void)observeValueForKeyPath:(NSString *)keyPath
+                      ofObject:(id)object
+                        change:(NSDictionary *)change
+                       context:(void *)context {
+    if ([keyPath isEqualToString:@"applicationIconBadgeNumber"]) {
+        [self updateTabBarBadgeNumber];
+    }
+}
 
 #pragma mark - Custom methods
 
+- (void)updateTabBarBadgeNumber {
+    UIApplication *application = [UIApplication sharedApplication];
+    NSInteger count = application.applicationIconBadgeNumber;
+
+    NSString *countString = count == 0 ? nil : [NSString stringWithFormat:@"%d", count];
+    self.navigationController.tabBarItem.badgeValue = countString;
+}
+
 - (void)refreshUnreadNotes {
     [Note refreshUnreadNotesWithContext:self.resultsController.managedObjectContext];
 }