[Sync] Refactor ChromeInvalidationClient and InvalidationNotifier

Use InvalidationClient's session state to determine whether notifications
are on/off, instead of just the XMPP connection state.

Combine ChromeInvalidationClient::Listener's invalidation methods into one.

Push the cached registered types state from InvalidationNotifier into
ChromeInvaildationClient.

BUG=
TEST=


Review URL: http://codereview.chromium.org/6820022

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@81152 0039d316-1c4b-4281-b951-d872f2087c98

6 files changed, 165 insertions, 134 deletions

diff --git a/chrome/browser/sync/notifier/chrome_invalidation_client.cc b/chrome/browser/sync/notifier/chrome_invalidation_client.cc
index cbc4782..074e481 100644
--- a/chrome/browser/sync/notifier/chrome_invalidation_client.cc
+++ b/chrome/browser/sync/notifier/chrome_invalidation_client.cc
@@ -75,6 +75,7 @@ void ChromeInvalidationClient::Start(
   ChangeBaseTask(base_task);
   registration_manager_.reset(
       new RegistrationManager(invalidation_client_.get()));
+  registration_manager_->SetRegisteredTypes(registered_types_);
 }
 
 void ChromeInvalidationClient::ChangeBaseTask(
@@ -105,8 +106,11 @@ void ChromeInvalidationClient::Stop() {
 void ChromeInvalidationClient::RegisterTypes(
     const syncable::ModelTypeSet& types) {
   DCHECK(non_thread_safe_.CalledOnValidThread());
-
-  registration_manager_->SetRegisteredTypes(types);
+  registered_types_ = types;
+  if (registration_manager_.get()) {
+    registration_manager_->SetRegisteredTypes(registered_types_);
+  }
+  // TODO(akalin): Clear invalidation versions for unregistered types.
 }
 
 void ChromeInvalidationClient::Invalidate(
@@ -119,12 +123,20 @@ void ChromeInvalidationClient::Invalidate(
   if (!ObjectIdToRealModelType(invalidation.object_id(), &model_type)) {
     LOG(WARNING) << "Could not get invalidation model type; "
                  << "invalidating everything";
-    listener_->OnInvalidateAll();
+    EmitInvalidation(registered_types_, std::string());
     RunAndDeleteClosure(callback);
     return;
   }
-  // TODO(akalin): Consider moving this version-checking code to the
-  // invalidation client.
+  // The invalidation API spec allows for the possibility of redundant
+  // invalidations, so keep track of the max versions and drop
+  // invalidations with old versions.
+  //
+  // TODO(akalin): Now that we keep track of registered types, we
+  // should drop invalidations for unregistered types.  We may also
+  // have to filter it at a higher level, as invalidations for
+  // newly-unregistered types may already be in flight.
+  //
+  // TODO(akalin): Persist |max_invalidation_versions_| somehow.
   if (invalidation.version() != UNKNOWN_OBJECT_VERSION) {
     std::map<syncable::ModelType, int64>::const_iterator it =
         max_invalidation_versions_.find(model_type);
@@ -136,12 +148,15 @@ void ChromeInvalidationClient::Invalidate(
     }
     max_invalidation_versions_[model_type] = invalidation.version();
   }
+
   std::string payload;
   // payload() CHECK()'s has_payload(), so we must check it ourselves first.
   if (invalidation.has_payload())
     payload = invalidation.payload();
 
-  listener_->OnInvalidate(model_type, payload);
+  syncable::ModelTypeSet types;
+  types.insert(model_type);
+  EmitInvalidation(types, payload);
   // TODO(akalin): We should really |callback| only after we get the
   // updates from the sync server. (see http://crbug.com/78462).
   RunAndDeleteClosure(callback);
@@ -154,12 +169,23 @@ void ChromeInvalidationClient::InvalidateAll(
   DCHECK(non_thread_safe_.CalledOnValidThread());
   DCHECK(invalidation::IsCallbackRepeatable(callback));
   VLOG(1) << "InvalidateAll";
-  listener_->OnInvalidateAll();
+  EmitInvalidation(registered_types_, std::string());
   // TODO(akalin): We should really |callback| only after we get the
   // updates from the sync server. (see http://crbug.com/76482).
   RunAndDeleteClosure(callback);
 }
 
+void ChromeInvalidationClient::EmitInvalidation(
+    const syncable::ModelTypeSet& types, const std::string& payload) {
+  DCHECK(non_thread_safe_.CalledOnValidThread());
+  // TODO(akalin): Move all uses of ModelTypeBitSet for invalidations
+  // to ModelTypeSet.
+  syncable::ModelTypePayloadMap type_payloads =
+      syncable::ModelTypePayloadMapFromBitSet(
+          syncable::ModelTypeBitSetFromSet(types), payload);
+  listener_->OnInvalidate(type_payloads);
+}
+
 void ChromeInvalidationClient::RegistrationStateChanged(
     const invalidation::ObjectId& object_id,
     invalidation::RegistrationState new_state,
@@ -194,6 +220,11 @@ void ChromeInvalidationClient::AllRegistrationsLost(
   RunAndDeleteClosure(callback);
 }
 
+void ChromeInvalidationClient::SessionStatusChanged(bool has_session) {
+  VLOG(1) << "SessionStatusChanged: " << has_session;
+  listener_->OnSessionStatusChanged(has_session);
+}
+
 void ChromeInvalidationClient::WriteState(const std::string& state) {
   DCHECK(non_thread_safe_.CalledOnValidThread());
   CHECK(state_writer_);
diff --git a/chrome/browser/sync/notifier/chrome_invalidation_client.h b/chrome/browser/sync/notifier/chrome_invalidation_client.h
index c38d88c..7ab0f3a 100644
--- a/chrome/browser/sync/notifier/chrome_invalidation_client.h
+++ b/chrome/browser/sync/notifier/chrome_invalidation_client.h
@@ -13,6 +13,7 @@
 #include <string>
 
 #include "base/basictypes.h"
+#include "base/compiler_specific.h"
 #include "base/memory/scoped_callback_factory.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/memory/weak_ptr.h"
@@ -20,8 +21,13 @@
 #include "chrome/browser/sync/notifier/chrome_system_resources.h"
 #include "chrome/browser/sync/notifier/state_writer.h"
 #include "chrome/browser/sync/syncable/model_type.h"
+#include "chrome/browser/sync/syncable/model_type_payload_map.h"
 #include "google/cacheinvalidation/invalidation-client.h"
 
+// TODO(akalin): Move invalidation::InvalidationListener into its own
+// file and include that instead of invalidation-client.h (which
+// includes generated protobuf header files).
+
 namespace talk_base {
 class Task;
 }  // namespace
@@ -31,9 +37,6 @@ namespace sync_notifier {
 class CacheInvalidationPacketHandler;
 class RegistrationManager;
 
-// TODO(akalin): Hook this up to a NetworkChangeNotifier so we can
-// properly notify invalidation_client_.
-
 class ChromeInvalidationClient
     : public invalidation::InvalidationListener,
       public StateWriter {
@@ -42,12 +45,10 @@ class ChromeInvalidationClient
    public:
     virtual ~Listener();
 
-    virtual void OnInvalidate(syncable::ModelType model_type,
-                              const std::string& payload) = 0;
+    virtual void OnInvalidate(
+        const syncable::ModelTypePayloadMap& type_payloads) = 0;
 
-    // TODO(akalin): Remove OnInvalidateAll() and have OnInvalidate
-    // take a ModelTypePayloadMap.
-    virtual void OnInvalidateAll() = 0;
+    virtual void OnSessionStatusChanged(bool has_session) = 0;
   };
 
   ChromeInvalidationClient();
@@ -70,28 +71,22 @@ class ChromeInvalidationClient
   void ChangeBaseTask(base::WeakPtr<talk_base::Task> base_task);
 
   // Register the sync types that we're interested in getting
-  // notifications for.  Must only be called between calls to Start()
-  // and Stop().
+  // notifications for.  May be called at any time.
   void RegisterTypes(const syncable::ModelTypeSet& types);
 
   // invalidation::InvalidationListener implementation.
-  //
-  // TODO(akalin): Move these into a private inner class.
-
   virtual void Invalidate(const invalidation::Invalidation& invalidation,
-                          invalidation::Closure* callback);
-
-  virtual void InvalidateAll(invalidation::Closure* callback);
-
+                          invalidation::Closure* callback) OVERRIDE;
+  virtual void InvalidateAll(invalidation::Closure* callback) OVERRIDE;
   virtual void RegistrationStateChanged(
       const invalidation::ObjectId& object_id,
       invalidation::RegistrationState new_state,
-      const invalidation::UnknownHint& unknown_hint);
-
-  virtual void AllRegistrationsLost(invalidation::Closure* callback);
+      const invalidation::UnknownHint& unknown_hint) OVERRIDE;
+  virtual void AllRegistrationsLost(invalidation::Closure* callback) OVERRIDE;
+  virtual void SessionStatusChanged(bool has_session) OVERRIDE;
 
   // StateWriter implementation.
-  virtual void WriteState(const std::string& state);
+  virtual void WriteState(const std::string& state) OVERRIDE;
 
  private:
   friend class ChromeInvalidationClientTest;
@@ -99,6 +94,8 @@ class ChromeInvalidationClient
   // Should only be called between calls to Start() and Stop().
   void HandleOutboundPacket(
       invalidation::NetworkEndpoint* const& network_endpoint);
+  void EmitInvalidation(
+      const syncable::ModelTypeSet& types, const std::string& payload);
 
   base::NonThreadSafe non_thread_safe_;
   ChromeSystemResources chrome_system_resources_;
@@ -112,6 +109,8 @@ class ChromeInvalidationClient
       cache_invalidation_packet_handler_;
   scoped_ptr<RegistrationManager> registration_manager_;
   std::map<syncable::ModelType, int64> max_invalidation_versions_;
+  // Stored to pass to |registration_manager_| on start.
+  syncable::ModelTypeSet registered_types_;
 
   DISALLOW_COPY_AND_ASSIGN(ChromeInvalidationClient);
 };
diff --git a/chrome/browser/sync/notifier/chrome_invalidation_client_unittest.cc b/chrome/browser/sync/notifier/chrome_invalidation_client_unittest.cc
index 50d9bbf..5a99e77 100644
--- a/chrome/browser/sync/notifier/chrome_invalidation_client_unittest.cc
+++ b/chrome/browser/sync/notifier/chrome_invalidation_client_unittest.cc
@@ -8,6 +8,7 @@
 #include "chrome/browser/sync/notifier/chrome_invalidation_client.h"
 #include "chrome/browser/sync/notifier/state_writer.h"
 #include "chrome/browser/sync/syncable/model_type.h"
+#include "chrome/browser/sync/syncable/model_type_payload_map.h"
 #include "jingle/notifier/base/fake_base_task.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -29,9 +30,8 @@ static const int64 kUnknownVersion =
 
 class MockListener : public ChromeInvalidationClient::Listener {
  public:
-  MOCK_METHOD2(OnInvalidate, void(syncable::ModelType,
-                                  const std::string& payload));
-  MOCK_METHOD0(OnInvalidateAll, void());
+  MOCK_METHOD1(OnInvalidate, void(const syncable::ModelTypePayloadMap&));
+  MOCK_METHOD1(OnSessionStatusChanged, void(bool));
 };
 
 class MockStateWriter : public StateWriter {
@@ -103,25 +103,53 @@ class ChromeInvalidationClientTest : public testing::Test {
   ChromeInvalidationClient client_;
 };
 
+namespace {
+
+syncable::ModelTypePayloadMap MakeMap(syncable::ModelType model_type,
+                                      const std::string& payload) {
+  syncable::ModelTypePayloadMap type_payloads;
+  type_payloads[model_type] = payload;
+  return type_payloads;
+}
+
+syncable::ModelTypePayloadMap MakeMapFromSet(syncable::ModelTypeSet types,
+                                             const std::string& payload) {
+  syncable::ModelTypePayloadMap type_payloads;
+  for (syncable::ModelTypeSet::const_iterator it = types.begin();
+       it != types.end(); ++it) {
+    type_payloads[*it] = payload;
+  }
+  return type_payloads;
+}
+
+}  // namespace
+
 TEST_F(ChromeInvalidationClientTest, InvalidateBadObjectId) {
-  EXPECT_CALL(mock_listener_, OnInvalidateAll());
+  syncable::ModelTypeSet types;
+  types.insert(syncable::BOOKMARKS);
+  types.insert(syncable::APPS);
+  client_.RegisterTypes(types);
+  EXPECT_CALL(mock_listener_, OnInvalidate(MakeMapFromSet(types, "")));
   FireInvalidate("bad", 1, NULL);
 }
 
 TEST_F(ChromeInvalidationClientTest, InvalidateNoPayload) {
-  EXPECT_CALL(mock_listener_, OnInvalidate(syncable::BOOKMARKS, ""));
+  EXPECT_CALL(mock_listener_,
+              OnInvalidate(MakeMap(syncable::BOOKMARKS, "")));
   FireInvalidate("BOOKMARK", 1, NULL);
 }
 
 TEST_F(ChromeInvalidationClientTest, InvalidateWithPayload) {
-  EXPECT_CALL(mock_listener_, OnInvalidate(syncable::PREFERENCES, "payload"));
+  EXPECT_CALL(mock_listener_,
+              OnInvalidate(MakeMap(syncable::PREFERENCES, "payload")));
   FireInvalidate("PREFERENCE", 1, "payload");
 }
 
 TEST_F(ChromeInvalidationClientTest, InvalidateVersion) {
   using ::testing::Mock;
 
-  EXPECT_CALL(mock_listener_, OnInvalidate(syncable::APPS, ""));
+  EXPECT_CALL(mock_listener_,
+              OnInvalidate(MakeMap(syncable::APPS, "")));
 
   // Should trigger.
   FireInvalidate("APP", 1, NULL);
@@ -133,7 +161,8 @@ TEST_F(ChromeInvalidationClientTest, InvalidateVersion) {
 }
 
 TEST_F(ChromeInvalidationClientTest, InvalidateUnknownVersion) {
-  EXPECT_CALL(mock_listener_, OnInvalidate(syncable::EXTENSIONS, ""))
+  EXPECT_CALL(mock_listener_,
+              OnInvalidate(MakeMap(syncable::EXTENSIONS, "")))
       .Times(2);
 
   // Should trigger twice.
@@ -144,8 +173,15 @@ TEST_F(ChromeInvalidationClientTest, InvalidateUnknownVersion) {
 TEST_F(ChromeInvalidationClientTest, InvalidateVersionMultipleTypes) {
   using ::testing::Mock;
 
-  EXPECT_CALL(mock_listener_, OnInvalidate(syncable::APPS, ""));
-  EXPECT_CALL(mock_listener_, OnInvalidate(syncable::EXTENSIONS, ""));
+  syncable::ModelTypeSet types;
+  types.insert(syncable::BOOKMARKS);
+  types.insert(syncable::APPS);
+  client_.RegisterTypes(types);
+
+  EXPECT_CALL(mock_listener_,
+              OnInvalidate(MakeMap(syncable::APPS, "")));
+  EXPECT_CALL(mock_listener_,
+              OnInvalidate(MakeMap(syncable::EXTENSIONS, "")));
 
   // Should trigger both.
   FireInvalidate("APP", 3, NULL);
@@ -159,9 +195,18 @@ TEST_F(ChromeInvalidationClientTest, InvalidateVersionMultipleTypes) {
 
   Mock::VerifyAndClearExpectations(&mock_listener_);
 
-  EXPECT_CALL(mock_listener_, OnInvalidate(syncable::PREFERENCES, ""));
-  EXPECT_CALL(mock_listener_, OnInvalidate(syncable::EXTENSIONS, ""));
-  EXPECT_CALL(mock_listener_, OnInvalidate(syncable::APPS, ""));
+  // InvalidateAll shouldn't change any version state.
+  EXPECT_CALL(mock_listener_, OnInvalidate(MakeMapFromSet(types, "")));
+  FireInvalidateAll();
+
+  Mock::VerifyAndClearExpectations(&mock_listener_);
+
+  EXPECT_CALL(mock_listener_,
+              OnInvalidate(MakeMap(syncable::PREFERENCES, "")));
+  EXPECT_CALL(mock_listener_,
+              OnInvalidate(MakeMap(syncable::EXTENSIONS, "")));
+  EXPECT_CALL(mock_listener_,
+              OnInvalidate(MakeMap(syncable::APPS, "")));
 
   // Should trigger all three.
   FireInvalidate("PREFERENCE", 5, NULL);
@@ -170,7 +215,23 @@ TEST_F(ChromeInvalidationClientTest, InvalidateVersionMultipleTypes) {
 }
 
 TEST_F(ChromeInvalidationClientTest, InvalidateAll) {
-  EXPECT_CALL(mock_listener_, OnInvalidateAll());
+  syncable::ModelTypeSet types;
+  types.insert(syncable::PREFERENCES);
+  types.insert(syncable::EXTENSIONS);
+  client_.RegisterTypes(types);
+  EXPECT_CALL(mock_listener_, OnInvalidate(MakeMapFromSet(types, "")));
+  FireInvalidateAll();
+}
+
+TEST_F(ChromeInvalidationClientTest, RegisterTypes) {
+  syncable::ModelTypeSet types;
+  types.insert(syncable::PREFERENCES);
+  types.insert(syncable::EXTENSIONS);
+  client_.RegisterTypes(types);
+  // Registered types should be preserved across Stop/Start.
+  TearDown();
+  SetUp();
+  EXPECT_CALL(mock_listener_, OnInvalidate(MakeMapFromSet(types, "")));
   FireInvalidateAll();
 }
 
diff --git a/chrome/browser/sync/notifier/invalidation_notifier.cc b/chrome/browser/sync/notifier/invalidation_notifier.cc
index 9dd6a24..a732a10 100644
--- a/chrome/browser/sync/notifier/invalidation_notifier.cc
+++ b/chrome/browser/sync/notifier/invalidation_notifier.cc
@@ -83,12 +83,7 @@ void InvalidationNotifier::UpdateCredentials(
 void InvalidationNotifier::UpdateEnabledTypes(
     const syncable::ModelTypeSet& types) {
   DCHECK(non_thread_safe_.CalledOnValidThread());
-  enabled_types_ = types;
-  if (state_ >= STARTED) {
-    invalidation_client_.RegisterTypes(enabled_types_);
-  }
-  // If |invalidation_client_| hasn't been started yet, it will
-  // register |enabled_types_| in OnConnect().
+  invalidation_client_.RegisterTypes(types);
 }
 
 void InvalidationNotifier::SendNotification() {
@@ -111,35 +106,25 @@ void InvalidationNotifier::OnConnect(
     invalidation_client_.Start(
         kClientId, client_info_, invalidation_state_, this, this, base_task);
     invalidation_state_.clear();
-    invalidation_client_.RegisterTypes(enabled_types_);
     state_ = STARTED;
   }
-  FOR_EACH_OBSERVER(SyncNotifierObserver, observers_,
-                    OnNotificationStateChange(true));
 }
 
 void InvalidationNotifier::OnDisconnect() {
   DCHECK(non_thread_safe_.CalledOnValidThread());
-  FOR_EACH_OBSERVER(SyncNotifierObserver, observers_,
-                    OnNotificationStateChange(false));
+  VLOG(1) << "OnDisconnect";
 }
 
 void InvalidationNotifier::OnInvalidate(
-    syncable::ModelType model_type, const std::string& payload) {
+    const syncable::ModelTypePayloadMap& type_payloads) {
   DCHECK(non_thread_safe_.CalledOnValidThread());
-  DCHECK_GE(model_type, syncable::FIRST_REAL_MODEL_TYPE);
-  DCHECK_LT(model_type, syncable::MODEL_TYPE_COUNT);
-  VLOG(1) << "OnInvalidate: " << syncable::ModelTypeToString(model_type)
-          << " " << payload;
-  syncable::ModelTypeSet types;
-  types.insert(model_type);
-  EmitInvalidation(types, payload);
+  FOR_EACH_OBSERVER(SyncNotifierObserver, observers_,
+                    OnIncomingNotification(type_payloads));
 }
 
-void InvalidationNotifier::OnInvalidateAll() {
-  DCHECK(non_thread_safe_.CalledOnValidThread());
-  VLOG(1) << "OnInvalidateAll";
-  EmitInvalidation(enabled_types_, std::string());
+void InvalidationNotifier::OnSessionStatusChanged(bool has_session) {
+  FOR_EACH_OBSERVER(SyncNotifierObserver, observers_,
+                    OnNotificationStateChange(has_session));
 }
 
 void InvalidationNotifier::WriteState(const std::string& state) {
@@ -148,16 +133,4 @@ void InvalidationNotifier::WriteState(const std::string& state) {
   FOR_EACH_OBSERVER(SyncNotifierObserver, observers_, StoreState(state));
 }
 
-void InvalidationNotifier::EmitInvalidation(
-    const syncable::ModelTypeSet& types, const std::string& payload) {
-  DCHECK(non_thread_safe_.CalledOnValidThread());
-  // TODO(akalin): Move all uses of ModelTypeBitSet for invalidations
-  // to ModelTypeSet.
-  syncable::ModelTypePayloadMap type_payloads =
-      syncable::ModelTypePayloadMapFromBitSet(
-          syncable::ModelTypeBitSetFromSet(types), payload);
-  FOR_EACH_OBSERVER(SyncNotifierObserver, observers_,
-                    OnIncomingNotification(type_payloads));
-}
-
 }  // namespace sync_notifier
diff --git a/chrome/browser/sync/notifier/invalidation_notifier.h b/chrome/browser/sync/notifier/invalidation_notifier.h
index 81e64b3..762aa93 100644
--- a/chrome/browser/sync/notifier/invalidation_notifier.h
+++ b/chrome/browser/sync/notifier/invalidation_notifier.h
@@ -16,6 +16,7 @@
 #include <string>
 
 #include "base/basictypes.h"
+#include "base/compiler_specific.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/observer_list.h"
 #include "base/threading/non_thread_safe.h"
@@ -42,30 +43,28 @@ class InvalidationNotifier
   virtual ~InvalidationNotifier();
 
   // SyncNotifier implementation.
-  virtual void AddObserver(SyncNotifierObserver* observer);
-  virtual void RemoveObserver(SyncNotifierObserver* observer);
-  virtual void SetState(const std::string& state);
+  virtual void AddObserver(SyncNotifierObserver* observer) OVERRIDE;
+  virtual void RemoveObserver(SyncNotifierObserver* observer) OVERRIDE;
+  virtual void SetState(const std::string& state) OVERRIDE;
   virtual void UpdateCredentials(
-      const std::string& email, const std::string& token);
-  virtual void UpdateEnabledTypes(const syncable::ModelTypeSet& types);
-  virtual void SendNotification();
+      const std::string& email, const std::string& token) OVERRIDE;
+  virtual void UpdateEnabledTypes(
+      const syncable::ModelTypeSet& types) OVERRIDE;
+  virtual void SendNotification() OVERRIDE;
 
   // notifier::LoginDelegate implementation.
-  virtual void OnConnect(base::WeakPtr<talk_base::Task> base_task);
-  virtual void OnDisconnect();
+  virtual void OnConnect(base::WeakPtr<talk_base::Task> base_task) OVERRIDE;
+  virtual void OnDisconnect() OVERRIDE;
 
   // ChromeInvalidationClient::Listener implementation.
-  virtual void OnInvalidate(syncable::ModelType model_type,
-                            const std::string& payload);
-  virtual void OnInvalidateAll();
+  virtual void OnInvalidate(
+      const syncable::ModelTypePayloadMap& type_payloads) OVERRIDE;
+  virtual void OnSessionStatusChanged(bool has_session) OVERRIDE;
 
   // StateWriter implementation.
-  virtual void WriteState(const std::string& state);
+  virtual void WriteState(const std::string& state) OVERRIDE;
 
  private:
-  void EmitInvalidation(const syncable::ModelTypeSet& types,
-                        const std::string& payload);
-
   base::NonThreadSafe non_thread_safe_;
 
   // We start off in the STOPPED state.  When we get our initial
@@ -97,10 +96,6 @@ class InvalidationNotifier
   // The invalidation client.
   ChromeInvalidationClient invalidation_client_;
 
-  // Passed to |invalidation_client_| and also used to synthesize
-  // notifications by OnInvalidateAll().
-  syncable::ModelTypeSet enabled_types_;
-
   DISALLOW_COPY_AND_ASSIGN(InvalidationNotifier);
 };
 
diff --git a/chrome/browser/sync/notifier/invalidation_notifier_unittest.cc b/chrome/browser/sync/notifier/invalidation_notifier_unittest.cc
index e6c7fce..bcfa9a4 100644
--- a/chrome/browser/sync/notifier/invalidation_notifier_unittest.cc
+++ b/chrome/browser/sync/notifier/invalidation_notifier_unittest.cc
@@ -56,59 +56,31 @@ class InvalidationNotifierTest : public testing::Test {
 TEST_F(InvalidationNotifierTest, Basic) {
   InSequence dummy;
 
-  syncable::ModelTypeSet types;
-  types.insert(syncable::BOOKMARKS);
-  types.insert(syncable::AUTOFILL);
+  syncable::ModelTypePayloadMap type_payloads;
+  type_payloads[syncable::PREFERENCES] = "payload";
+  type_payloads[syncable::BOOKMARKS] = "";
+  type_payloads[syncable::AUTOFILL] = "";
 
   EXPECT_CALL(mock_observer_, OnNotificationStateChange(true));
   EXPECT_CALL(mock_observer_, StoreState("new_fake_state"));
-  {
-    syncable::ModelTypePayloadMap type_payloads;
-    type_payloads[syncable::PREFERENCES] = "payload";
-    EXPECT_CALL(mock_observer_, OnIncomingNotification(type_payloads));
-  }
-  {
-    syncable::ModelTypePayloadMap type_payloads;
-    type_payloads[syncable::BOOKMARKS] = "";
-    type_payloads[syncable::AUTOFILL] = "";
-    EXPECT_CALL(mock_observer_, OnIncomingNotification(type_payloads));
-  }
+  EXPECT_CALL(mock_observer_, OnIncomingNotification(type_payloads));
   EXPECT_CALL(mock_observer_, OnNotificationStateChange(false));
 
   invalidation_notifier_->SetState("fake_state");
   invalidation_notifier_->UpdateCredentials("foo@bar.com", "fake_token");
-  invalidation_notifier_->UpdateEnabledTypes(types);
 
   invalidation_notifier_->OnConnect(fake_base_task_.AsWeakPtr());
+  invalidation_notifier_->OnSessionStatusChanged(true);
 
   invalidation_notifier_->WriteState("new_fake_state");
 
-  // Even though preferences isn't in the set of enabled types, we
-  // should still propagate the notification.
-  invalidation_notifier_->OnInvalidate(syncable::PREFERENCES, "payload");
-  invalidation_notifier_->OnInvalidateAll();
+  invalidation_notifier_->OnInvalidate(type_payloads);
 
+  // Shouldn't trigger notification state change.
   invalidation_notifier_->OnDisconnect();
-}
-
-TEST_F(InvalidationNotifierTest, UpdateEnabledTypes) {
-  InSequence dummy;
-
-  syncable::ModelTypeSet types;
-  types.insert(syncable::BOOKMARKS);
-  types.insert(syncable::AUTOFILL);
-
-  EXPECT_CALL(mock_observer_, OnNotificationStateChange(true));
-  {
-    syncable::ModelTypePayloadMap type_payloads;
-    type_payloads[syncable::BOOKMARKS] = "";
-    type_payloads[syncable::AUTOFILL] = "";
-    EXPECT_CALL(mock_observer_, OnIncomingNotification(type_payloads));
-  }
-
   invalidation_notifier_->OnConnect(fake_base_task_.AsWeakPtr());
-  invalidation_notifier_->UpdateEnabledTypes(types);
-  invalidation_notifier_->OnInvalidateAll();
+
+  invalidation_notifier_->OnSessionStatusChanged(false);
 }
 
 }  // namespace