CrOS Shelf: Use AppRegistryCache for shelf notification badging

Before this, the MessageCenterObserver was used to allow the ShelfModel
to keep track of existing notifications. This allowed the ShelfModel to
correctly update the apps that had a notification.

In this CL, the ShelfController no longer observes the MessageCenter and
instead observes the AppRegistryCache to get notification updates. With
this method, the ShelfModel no longer needs to keep a map of existing
notifications. Instead, the AppRegistryCache calls OnAppUpdate()
whenever an app's notification badge status has changed.

Also, in order to get the initial notification badge status for new
ShelfItems, the ShelfControlelr observes changes to the ShelfModel.
When a new item is added to the model, the ShelfController uses the
AppRegistryCache to update the model with the current notification
badge status.

Bug: 1080827
Change-Id: I1c37626170387a8e81cec0f6f3c06026ac62a976
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2219164Reviewed-by: Nancy Wang <nancylingwang@chromium.org>
Reviewed-by: Xiyuan Xia <xiyuan@chromium.org>
Commit-Queue: Matthew Mourgos <mmourgos@chromium.org>
Cr-Commit-Position: refs/heads/master@{#786304}

ash/BUILD.gn

@@ -1575,6 +1575,7 @@ component("ash") {
     "//components/pref_registry",
     "//components/prefs",
     "//components/quirks",
+    "//components/services/app_service/public/cpp:app_update",
     "//components/session_manager:base",
     "//components/strings",
     "//components/sync",
@@ -2245,6 +2246,7 @@ test("ash_unittests") {
     "//components/password_manager/core/browser:hash_password_manager",
     "//components/prefs:test_support",
     "//components/quirks",
+    "//components/services/app_service/public/cpp:app_update",
     "//components/sync_preferences:test_support",
     "//components/ukm:test_support",
     "//components/user_manager",

ash/DEPS

@@ -12,6 +12,7 @@ include_rules = [
   "+components/pref_registry",
   "+components/prefs",
   "+components/quirks",
+  "+components/services/app_service/public",
   "+components/session_manager",
   "+components/strings",
   "+components/sync",

ash/public/cpp/shelf_model.cc

@@ -118,8 +118,6 @@ int ShelfModel::AddAt(int index, const ShelfItem& item) {
   DCHECK_EQ(ItemIndexByID(item.id), -1) << " The id is not unique: " << item.id;
   index = ValidateInsertionIndex(item.type, index);
   items_.insert(items_.begin() + index, item);
-  items_[index].has_notification =
-      app_id_to_notification_id_.count(item.id.app_id) > 0;
   for (auto& observer : observers_)
     observer.ShelfItemAdded(index);
   return index;
@@ -254,53 +252,6 @@ void ShelfModel::OnItemReturnedFromRipOff(int index) {
     observer.ShelfItemReturnedFromRipOff(index);
 }
 
-void ShelfModel::RemoveNotificationRecord(const std::string& notification_id) {
-  auto notification_id_it = notification_id_to_app_id_.find(notification_id);
-
-  // Two maps are required here because when this notification has been
-  // delivered, the MessageCenter has already deleted the notification, so we
-  // can't fetch the corresponding App Id.
-  // If we have a record of this notification, erase it from both maps.
-  if (notification_id_it == notification_id_to_app_id_.end())
-    return;
-
-  // Save the AppId so the app can be updated.
-  const std::string app_id = notification_id_it->second;
-
-  auto app_id_it = app_id_to_notification_id_.find(app_id);
-
-  // Remove the notification_id.
-  app_id_it->second.erase(notification_id);
-
-  // If the set is empty erase the pair.
-  if (app_id_it->second.empty())
-    app_id_to_notification_id_.erase(app_id_it);
-
-  // Erase the pair in the NotificationId -> AppId map.
-  notification_id_to_app_id_.erase(notification_id_it);
-
-  UpdateItemNotificationsAndNotifyObservers(app_id);
-}
-
-void ShelfModel::AddNotificationRecord(const std::string& app_id,
-                                       const std::string& notification_id) {
-  auto it = app_id_to_notification_id_.find(app_id);
-  if (it != app_id_to_notification_id_.end()) {
-    // The app_id exists in the map, modify the set.
-    it->second.insert(notification_id);
-  } else {
-    // The app_id hasn't been recorded yet, create a set.
-    app_id_to_notification_id_.insert(
-        std::pair<std::string, std::set<std::string>>(app_id,
-                                                      {notification_id}));
-  }
-
-  notification_id_to_app_id_.insert(
-      std::pair<std::string, std::string>(notification_id, app_id));
-
-  UpdateItemNotificationsAndNotifyObservers(app_id);
-}
-
 int ShelfModel::ItemIndexByID(const ShelfID& shelf_id) const {
   ShelfItems::const_iterator i = ItemByID(shelf_id);
   return i == items_.end() ? -1 : static_cast<int>(i - items_.begin());
@@ -399,18 +350,17 @@ int ShelfModel::ValidateInsertionIndex(ShelfItemType type, int index) const {
   return index;
 }
 
-void ShelfModel::UpdateItemNotificationsAndNotifyObservers(
-    const std::string& app_id) {
+void ShelfModel::UpdateItemNotification(const std::string& app_id,
+                                        bool has_badge) {
   int index = ItemIndexByAppID(app_id);
   // If the item is not pinned or active on the shelf.
   if (index == -1)
     return;
 
-  const bool has_notification = app_id_to_notification_id_.count(app_id) > 0;
-  if (items_[index].has_notification == has_notification)
+  if (items_[index].has_notification == has_badge)
     return;
 
-  items_[index].has_notification = has_notification;
+  items_[index].has_notification = has_badge;
 
   for (auto& observer : observers_)
     observer.ShelfItemChanged(index, items_[index]);

ash/public/cpp/shelf_model.h

@@ -140,13 +140,9 @@ class ASH_PUBLIC_EXPORT ShelfModel {
   // dragged back onto the shelf (it is still being dragged).
   void OnItemReturnedFromRipOff(int index);
 
-  // Adds a record of the notification with this app id and notifies observers.
-  void AddNotificationRecord(const std::string& app_id,
-                             const std::string& notification_id);
-
-  // Removes the record of the notification with matching ID and notifies
-  // observers.
-  void RemoveNotificationRecord(const std::string& notification_id);
+  // Update the ShelfItem with |app_id| to set whether the item currently has a
+  // notification.
+  void UpdateItemNotification(const std::string& app_id, bool has_badge);
 
   // Returns the index of the item with id |shelf_id|, or -1 if none exists.
   int ItemIndexByID(const ShelfID& shelf_id) const;
@@ -192,10 +188,6 @@ class ASH_PUBLIC_EXPORT ShelfModel {
   // returns the new value.
   int ValidateInsertionIndex(ShelfItemType type, int index) const;
 
-  // Finds the app corresponding to |app_id|, sets ShelfItem.has_notification,
-  // and notifies observers.
-  void UpdateItemNotificationsAndNotifyObservers(const std::string& app_id);
-
   ShelfItems items_;
 
   // The shelf ID of the currently active shelf item, or an empty ID if
@@ -208,11 +200,6 @@ class ASH_PUBLIC_EXPORT ShelfModel {
   // user interaction.
   int current_mutation_is_user_triggered_ = 0;
 
-  // Maps one app id to a set of all matching notification ids.
-  std::map<std::string, std::set<std::string>> app_id_to_notification_id_;
-  // Maps one notification id to one app id.
-  std::map<std::string, std::string> notification_id_to_app_id_;
-
   base::ObserverList<ShelfModelObserver>::Unchecked observers_;
 
   std::map<ShelfID, std::unique_ptr<ShelfItemDelegate>>

ash/public/cpp/shelf_model_unittest.cc

@@ -411,7 +411,6 @@ TEST_F(ShelfModelTest, RunningAppPinning) {
 // Tests that apps are updated properly when notifications are added or removed.
 TEST_F(ShelfModelTest, AddRemoveNotification) {
   const std::string app_id("app_id");
-  const std::string notification_id("notification_id");
 
   // Add an example running app.
   ShelfItem item;
@@ -422,99 +421,12 @@ TEST_F(ShelfModelTest, AddRemoveNotification) {
 
   EXPECT_FALSE(model_->items()[index].has_notification);
 
-  // Add a notification for the app.
-  model_->AddNotificationRecord(app_id, notification_id);
-
-  EXPECT_TRUE(model_->items()[index].has_notification);
-
-  // Remove the notification.
-  model_->RemoveNotificationRecord(notification_id);
-
-  EXPECT_FALSE(model_->items()[index].has_notification);
-}
-
-// Tests that apps pick up their notifications when they are added.
-TEST_F(ShelfModelTest, AddAppAfterNotification) {
-  const std::string app_id("app_id");
-  const std::string notification_id("notification_id");
-
-  ShelfItem item;
-  item.type = TYPE_APP;
-  item.status = STATUS_RUNNING;
-  item.id = ShelfID(app_id);
-
-  // Add a notification for the app.
-  model_->AddNotificationRecord(app_id, notification_id);
-
-  // Add an app with a matching app id.
-  const int index = model_->Add(item);
-
-  EXPECT_TRUE(model_->items()[index].has_notification);
-
-  // Remove and re-add the app.
-  model_->RemoveItemAt(index);
-  EXPECT_EQ(index, model_->Add(item));
-
-  // Test that the notification persists.
-  EXPECT_TRUE(model_->items()[index].has_notification);
-}
-
-// Tests that pinned apps pick up their notifications if they were recieved
-// before the app existed on the shelf.
-TEST_F(ShelfModelTest, PinAppAfterNotification) {
-  const std::string app_id("app_id");
-  const std::string notification_id("notification_id");
-
-  // Add an example app, but don't pin it.
-  ShelfItem item;
-  item.type = TYPE_APP;
-  item.status = STATUS_RUNNING;
-  item.id = ShelfID(app_id);
-
-  // Add a notification for the app.
-  model_->AddNotificationRecord(app_id, notification_id);
-
-  // Pin the app after the notification posts.
-  model_->PinAppWithID(app_id);
-
-  const int index = model_->ItemIndexByAppID(app_id);
-  EXPECT_TRUE(model_->items()[index].has_notification);
-
-  // Un-pin and re-pin the app.
-  model_->UnpinAppWithID(app_id);
-  model_->PinAppWithID(app_id);
-  EXPECT_EQ(index, model_->ItemIndexByAppID(app_id));
-
-  // Test that the notification persists.
-  EXPECT_TRUE(model_->items()[index].has_notification);
-}
-
-// Tests that the ShelfItem.has_notification is set to false only when there are
-// 0 notifications.
-TEST_F(ShelfModelTest, MultipleNotificationsPerAppBasic) {
-  const std::string app_id("app_id");
-  const std::string notification_id_0("notification_id_0");
-  // Add an example app.
-  ShelfItem item;
-  item.type = TYPE_APP;
-  item.status = STATUS_RUNNING;
-  item.id = ShelfID(app_id);
-  const int index = model_->Add(item);
-  EXPECT_EQ(index, model_->ItemIndexByAppID(app_id));
-
-  // Add the first notification for this app.
-  model_->AddNotificationRecord(app_id, notification_id_0);
-  // Add the second notification on the same app.
-  const std::string notification_id_1("notification_id_1");
-  model_->AddNotificationRecord(app_id, notification_id_1);
-  EXPECT_TRUE(model_->items()[index].has_notification);
-
-  // Remove one notification.
-  model_->RemoveNotificationRecord(notification_id_1);
+  // Update to add a notification for the app.
+  model_->UpdateItemNotification(app_id, true /* has_badge */);
   EXPECT_TRUE(model_->items()[index].has_notification);
 
-  // Remove the last notification.
-  model_->RemoveNotificationRecord(notification_id_0);
+  // Update to remove the notification for the app.
+  model_->UpdateItemNotification(app_id, false /* has_badge */);
   EXPECT_FALSE(model_->items()[index].has_notification);
 }
 

ash/shelf/shelf_app_button.cc

@@ -515,6 +515,13 @@ bool ShelfAppButton::ShouldEnterPushedState(const ui::Event& event) {
 }
 
 void ShelfAppButton::ReflectItemStatus(const ShelfItem& item) {
+  if (features::IsNotificationIndicatorEnabled()) {
+    if (item.has_notification)
+      AddState(ShelfAppButton::STATE_NOTIFICATION);
+    else
+      ClearState(ShelfAppButton::STATE_NOTIFICATION);
+  }
+
   const ShelfID active_id = shelf_view_->model()->active_shelf_id();
   if (!active_id.IsNull() && item.id == active_id) {
     // The active status trumps all other statuses.
@@ -540,13 +547,6 @@ void ShelfAppButton::ReflectItemStatus(const ShelfItem& item) {
       AddState(ShelfAppButton::STATE_ATTENTION);
       break;
   }
-
-  if (features::IsNotificationIndicatorEnabled()) {
-    if (item.has_notification)
-      AddState(ShelfAppButton::STATE_NOTIFICATION);
-    else
-      ClearState(ShelfAppButton::STATE_NOTIFICATION);
-  }
 }
 
 bool ShelfAppButton::IsIconSizeCurrent() {

ash/shelf/shelf_controller.cc

@@ -21,6 +21,7 @@
 #include "components/prefs/pref_change_registrar.h"
 #include "components/prefs/pref_registry_simple.h"
 #include "components/prefs/pref_service.h"
+#include "components/services/app_service/public/cpp/app_registry_cache_wrapper.h"
 #include "ui/base/l10n/l10n_util.h"
 #include "ui/base/ui_base_features.h"
 #include "ui/display/display.h"
@@ -106,9 +107,7 @@ ShelfController::ShelfController()
   Shell::Get()->session_controller()->AddObserver(this);
   Shell::Get()->tablet_mode_controller()->AddObserver(this);
   Shell::Get()->window_tree_host_manager()->AddObserver(this);
-
-  if (is_notification_indicator_enabled_)
-    message_center_observer_.Add(message_center::MessageCenter::Get());
+  model_.AddObserver(this);
 }
 
 ShelfController::~ShelfController() {
@@ -116,6 +115,7 @@ ShelfController::~ShelfController() {
 }
 
 void ShelfController::Shutdown() {
+  model_.RemoveObserver(this);
   Shell::Get()->window_tree_host_manager()->RemoveObserver(this);
   Shell::Get()->tablet_mode_controller()->RemoveObserver(this);
   Shell::Get()->session_controller()->RemoveObserver(this);
@@ -149,6 +149,27 @@ void ShelfController::OnActiveUserPrefServiceChanged(
                               base::BindRepeating(&SetShelfAutoHideFromPrefs));
   pref_change_registrar_->Add(prefs::kShelfPreferences,
                               base::BindRepeating(&SetShelfBehaviorsFromPrefs));
+
+  if (is_notification_indicator_enabled_) {
+    // Observe AppRegistryCache for the current active account to get
+    // notification updates.
+    AccountId account_id =
+        Shell::Get()->session_controller()->GetActiveAccountId();
+    cache_ =
+        apps::AppRegistryCacheWrapper::Get().GetAppRegistryCache(account_id);
+
+    Observe(cache_);
+
+    if (cache_) {
+      // Update the notification badge indicator for all apps. This will also
+      // ensure that apps have the correct notification badge value for the
+      // multiprofile case when switching between users.
+      cache_->ForEachApp([this](const apps::AppUpdate& update) {
+        bool has_badge = update.HasBadge() == apps::mojom::OptionalBool::kTrue;
+        model_.UpdateItemNotification(update.AppId(), has_badge);
+      });
+    }
+  }
 }
 
 void ShelfController::OnTabletModeStarted() {
@@ -196,38 +217,29 @@ void ShelfController::OnDisplayConfigurationChanged() {
   UpdateShelfVisibility();
 }
 
-void ShelfController::OnNotificationAdded(const std::string& notification_id) {
-  if (!is_notification_indicator_enabled_)
-    return;
-
-  message_center::Notification* notification =
-      message_center::MessageCenter::Get()->FindVisibleNotificationById(
-          notification_id);
-
-  if (!notification)
-    return;
-
-  // Skip this if the notification shouldn't badge an app.
-  if (notification->notifier_id().type !=
-          message_center::NotifierType::APPLICATION &&
-      notification->notifier_id().type !=
-          message_center::NotifierType::ARC_APPLICATION) {
-    return;
+void ShelfController::OnAppUpdate(const apps::AppUpdate& update) {
+  if (update.HasBadgeChanged()) {
+    bool has_badge = update.HasBadge() == apps::mojom::OptionalBool::kTrue;
+    model_.UpdateItemNotification(update.AppId(), has_badge);
   }
+}
 
-  // Skip this if the notification doesn't have a valid app id.
-  if (notification->notifier_id().id == kDefaultArcNotifierId)
-    return;
-
-  model_.AddNotificationRecord(notification->notifier_id().id, notification_id);
+void ShelfController::OnAppRegistryCacheWillBeDestroyed(
+    apps::AppRegistryCache* cache) {
+  Observe(nullptr);
 }
 
-void ShelfController::OnNotificationRemoved(const std::string& notification_id,
-                                            bool by_user) {
-  if (!is_notification_indicator_enabled_)
+void ShelfController::ShelfItemAdded(int index) {
+  if (!is_notification_indicator_enabled_ || !cache_)
     return;
 
-  model_.RemoveNotificationRecord(notification_id);
+  auto app_id = model_.items()[index].id.app_id;
+
+  // Update the notification badge indicator for the newly added shelf item.
+  cache_->ForOneApp(app_id, [this](const apps::AppUpdate& update) {
+    bool has_badge = update.HasBadge() == apps::mojom::OptionalBool::kTrue;
+    model_.UpdateItemNotification(update.AppId(), has_badge);
+  });
 }
 
 }  // namespace ash

ash/shelf/shelf_controller.h

@@ -12,10 +12,11 @@
 #include "ash/display/window_tree_host_manager.h"
 #include "ash/public/cpp/session/session_observer.h"
 #include "ash/public/cpp/shelf_model.h"
+#include "ash/public/cpp/shelf_model_observer.h"
 #include "ash/public/cpp/tablet_mode_observer.h"
 #include "base/scoped_observer.h"
-#include "ui/message_center/message_center.h"
-#include "ui/message_center/message_center_observer.h"
+#include "components/services/app_service/public/cpp/app_registry_cache.h"
+#include "components/services/app_service/public/cpp/app_update.h"
 
 class PrefChangeRegistrar;
 class PrefRegistrySimple;
@@ -24,10 +25,11 @@ namespace ash {
 
 // ShelfController owns the ShelfModel and manages shelf preferences.
 // ChromeLauncherController and related classes largely manage the ShelfModel.
-class ASH_EXPORT ShelfController : public message_center::MessageCenterObserver,
-                                   public SessionObserver,
+class ASH_EXPORT ShelfController : public SessionObserver,
                                    public TabletModeObserver,
-                                   public WindowTreeHostManager::Observer {
+                                   public WindowTreeHostManager::Observer,
+                                   public apps::AppRegistryCache::Observer,
+                                   public ShelfModelObserver {
  public:
   ShelfController();
   ~ShelfController() override;
@@ -40,11 +42,6 @@ class ASH_EXPORT ShelfController : public message_center::MessageCenterObserver,
   ShelfModel* model() { return &model_; }
 
  private:
-  // message_center::MessageCenterObserver:
-  void OnNotificationAdded(const std::string& notification_id) override;
-  void OnNotificationRemoved(const std::string& notification_id,
-                             bool by_user) override;
-
   // SessionObserver:
   void OnActiveUserPrefServiceChanged(PrefService* pref_service) override;
 
@@ -55,19 +52,27 @@ class ASH_EXPORT ShelfController : public message_center::MessageCenterObserver,
   // WindowTreeHostManager::Observer:
   void OnDisplayConfigurationChanged() override;
 
+  // apps::AppRegistryCache::Observer:
+  void OnAppUpdate(const apps::AppUpdate& update) override;
+  void OnAppRegistryCacheWillBeDestroyed(
+      apps::AppRegistryCache* cache) override;
+
+  // ShelfModelObserver:
+  void ShelfItemAdded(int index) override;
+
   // The shelf model shared by all shelf instances.
   ShelfModel model_;
 
   // Whether notification indicators are enabled for app icons in the shelf.
   const bool is_notification_indicator_enabled_;
 
-  ScopedObserver<message_center::MessageCenter,
-                 message_center::MessageCenterObserver>
-      message_center_observer_{this};
-
   // Observes user profile prefs for the shelf.
   std::unique_ptr<PrefChangeRegistrar> pref_change_registrar_;
 
+  // Observed to update notification badging on shelf items. Also used to get
+  // initial notification badge information when shelf items are added.
+  apps::AppRegistryCache* cache_ = nullptr;
+
   DISALLOW_COPY_AND_ASSIGN(ShelfController);
 };
 

ash/shelf/shelf_controller_unittest.cc

@@ -37,20 +37,7 @@ Shelf* GetShelfForDisplay(int64_t display_id) {
   return Shell::GetRootWindowControllerWithDisplayId(display_id)->shelf();
 }
 
-void BuildAndSendNotification(message_center::MessageCenter* message_center,
-                              const std::string& app_id,
-                              const std::string& notification_id) {
-  const message_center::NotifierId notifier_id(
-      message_center::NotifierType::APPLICATION, app_id);
-  std::unique_ptr<message_center::Notification> notification =
-      std::make_unique<message_center::Notification>(
-          message_center::NOTIFICATION_TYPE_SIMPLE, notification_id,
-          base::ASCIIToUTF16("Test Web Notification"),
-          base::ASCIIToUTF16("Notification message body."), gfx::Image(),
-          base::ASCIIToUTF16("www.test.org"), GURL(), notifier_id,
-          message_center::RichNotificationData(), nullptr /* delegate */);
-  message_center->AddNotification(std::move(notification));
-}
+}  // namespace
 
 using ShelfControllerTest = AshTestBase;
 
@@ -93,16 +80,34 @@ class ShelfControllerNotificationIndicatorTest : public AshTestBase {
     scoped_feature_list_.InitWithFeatures({features::kNotificationIndicator},
                                           {});
     AshTestBase::SetUp();
+
+    account_id_ = AccountId::FromUserEmail("test@gmail.com");
+  }
+
+  void SendAppUpdate(bool app_has_badge) {
+    ShelfController* controller = Shell::Get()->shelf_controller();
+
+    apps::mojom::App test_app;
+    test_app.app_id = "app_id";
+    if (app_has_badge)
+      test_app.has_badge = apps::mojom::OptionalBool::kTrue;
+    else
+      test_app.has_badge = apps::mojom::OptionalBool::kFalse;
+
+    apps::AppUpdate test_update(nullptr, &test_app /* delta */, account_id_);
+    static_cast<apps::AppRegistryCache::Observer*>(controller)
+        ->OnAppUpdate(test_update);
   }
 
  private:
+  AccountId account_id_;
   base::test::ScopedFeatureList scoped_feature_list_;
 
   DISALLOW_COPY_AND_ASSIGN(ShelfControllerNotificationIndicatorTest);
 };
 
-// Tests that the ShelfController keeps the ShelfModel updated on new
-// notifications.
+// Tests that the ShelfController keeps the ShelfModel updated on calls to
+// OnAppUpdate().
 TEST_F(ShelfControllerNotificationIndicatorTest, HasNotificationBasic) {
   ShelfController* controller = Shell::Get()->shelf_controller();
   const std::string app_id("app_id");
@@ -112,21 +117,13 @@ TEST_F(ShelfControllerNotificationIndicatorTest, HasNotificationBasic) {
   const int index = controller->model()->Add(item);
   EXPECT_FALSE(controller->model()->items()[index].has_notification);
 
-  // Add a notification for |item|.
-  message_center::MessageCenter* message_center =
-      message_center::MessageCenter::Get();
-  const std::string notification_id("notification_id");
-  BuildAndSendNotification(message_center, app_id, notification_id);
+  // Send an app update to ShelfController for adding a notification badge.
+  SendAppUpdate(true /* app_has_badge */);
 
   EXPECT_TRUE(controller->model()->items()[index].has_notification);
 
-  // Remove the app and pin it, the notification should persist.
-  controller->model()->RemoveItemAt(index);
-  controller->model()->PinAppWithID(app_id);
-
-  EXPECT_TRUE(controller->model()->items()[index].has_notification);
-
-  message_center->RemoveNotification(notification_id, true);
+  // Send an app update to ShelfController for removing a notification badge.
+  SendAppUpdate(false /* app_has_badge */);
 
   EXPECT_FALSE(controller->model()->items()[index].has_notification);
 }
@@ -386,5 +383,4 @@ TEST_F(ShelfControllerAppModeTest, AutoHideBehavior) {
   EXPECT_EQ(ShelfAutoHideBehavior::kAlwaysHidden, shelf->auto_hide_behavior());
 }
 
-}  // namespace
 }  // namespace ash

ash/shelf/shelf_view_unittest.cc

@@ -2110,73 +2110,26 @@ class NotificationIndicatorTest : public ShelfViewTest {
   DISALLOW_COPY_AND_ASSIGN(NotificationIndicatorTest);
 };
 
-// Tests that an item has a notification indicator when it recieves a
-// notification.
-TEST_F(NotificationIndicatorTest, AddedItemHasNotificationIndicator) {
-  const ShelfID id_0 = AddApp();
-  const std::string notification_id_0("notification_id_0");
-  const ShelfAppButton* button_0 = GetButtonByID(id_0);
+// Tests that an item has a notification badge indicator when the notification
+// is added and removed.
+TEST_F(NotificationIndicatorTest, ItemHasCorrectNotificationBadgeIndicator) {
+  const ShelfID item_id = AddApp();
+  const ShelfAppButton* shelf_app_button = GetButtonByID(item_id);
 
-  EXPECT_FALSE(GetItemByID(id_0).has_notification);
-  EXPECT_FALSE(button_0->state() & ShelfAppButton::STATE_NOTIFICATION);
+  EXPECT_FALSE(GetItemByID(item_id).has_notification);
+  EXPECT_FALSE(shelf_app_button->state() & ShelfAppButton::STATE_NOTIFICATION);
 
-  // Post a test notification after the item was added.
-  model_->AddNotificationRecord(id_0.app_id, notification_id_0);
+  // Add a notification for the new shelf item.
+  model_->UpdateItemNotification(item_id.app_id, true /* has_badge */);
 
-  EXPECT_TRUE(GetItemByID(id_0).has_notification);
-  EXPECT_TRUE(button_0->state() & ShelfAppButton::STATE_NOTIFICATION);
+  EXPECT_TRUE(GetItemByID(item_id).has_notification);
+  EXPECT_TRUE(shelf_app_button->state() & ShelfAppButton::STATE_NOTIFICATION);
 
-  // Post another notification for a non existing item.
-  const std::string next_app_id(GetNextAppId());
-  const std::string notification_id_1("notification_id_1");
-  model_->AddNotificationRecord(next_app_id, notification_id_1);
+  // Remove notification.
+  model_->UpdateItemNotification(item_id.app_id, false /* has_badge */);
 
-  // Add an item with matching app id.
-  const ShelfID id_1 = AddApp();
-
-  // Ensure that the app id assigned to |id_1| is the same as |next_app_id|.
-  EXPECT_EQ(next_app_id, id_1.app_id);
-  const ShelfAppButton* button_1 = GetButtonByID(id_1);
-  EXPECT_TRUE(GetItemByID(id_1).has_notification);
-  EXPECT_TRUE(button_1->state() & ShelfAppButton::STATE_NOTIFICATION);
-
-  // Remove all notifications.
-  model_->RemoveNotificationRecord(notification_id_0);
-  model_->RemoveNotificationRecord(notification_id_1);
-
-  EXPECT_FALSE(GetItemByID(id_0).has_notification);
-  EXPECT_FALSE(button_0->state() & ShelfAppButton::STATE_NOTIFICATION);
-  EXPECT_FALSE(GetItemByID(id_1).has_notification);
-  EXPECT_FALSE(button_1->state() & ShelfAppButton::STATE_NOTIFICATION);
-}
-
-// Tests that the notification indicator is active until all notifications have
-// been removed.
-TEST_F(NotificationIndicatorTest,
-       NotificationIndicatorStaysActiveUntilNotificationsAreGone) {
-  const ShelfID app = AddApp();
-  const ShelfAppButton* button = GetButtonByID(app);
-
-  // Add two notifications for the same app.
-  const std::string notification_id_0("notification_id_0");
-  model_->AddNotificationRecord(app.app_id, notification_id_0);
-  const std::string notification_id_1("notification_id_1");
-  model_->AddNotificationRecord(app.app_id, notification_id_1);
-
-  EXPECT_TRUE(GetItemByID(app).has_notification);
-  EXPECT_TRUE(button->state() & ShelfAppButton::STATE_NOTIFICATION);
-
-  // Remove one notification, indicator should stay active.
-  model_->RemoveNotificationRecord(notification_id_0);
-
-  EXPECT_TRUE(GetItemByID(app).has_notification);
-  EXPECT_TRUE(button->state() & ShelfAppButton::STATE_NOTIFICATION);
-
-  // Remove the last notification, indicator should not be active.
-  model_->RemoveNotificationRecord(notification_id_1);
-
-  EXPECT_FALSE(GetItemByID(app).has_notification);
-  EXPECT_FALSE(button->state() & ShelfAppButton::STATE_NOTIFICATION);
+  EXPECT_FALSE(GetItemByID(item_id).has_notification);
+  EXPECT_FALSE(shelf_app_button->state() & ShelfAppButton::STATE_NOTIFICATION);
 }
 
 class ShelfViewVisibleBoundsTest : public ShelfViewTest,