Fix crash when affiliated invalidation service becomes available

chrome/browser/chromeos/policy/affiliated_invalidation_service_provider_impl_unittest.cc

Index: chrome/browser/chromeos/policy/affiliated_invalidation_service_provider_impl_unittest.cc
diff --git a/chrome/browser/chromeos/policy/affiliated_invalidation_service_provider_impl_unittest.cc b/chrome/browser/chromeos/policy/affiliated_invalidation_service_provider_impl_unittest.cc
index 370d4a716a6e46cbff2cbfbda76cad17c9b967eb..bcbf68527f089ee3f9c57f4dbaa7e2bb514f110f 100644
--- a/chrome/browser/chromeos/policy/affiliated_invalidation_service_provider_impl_unittest.cc
+++ b/chrome/browser/chromeos/policy/affiliated_invalidation_service_provider_impl_unittest.cc
@@ -20,6 +20,7 @@
 #include "chrome/test/base/testing_profile_manager.h"
 #include "chromeos/cryptohome/system_salt_getter.h"
 #include "chromeos/dbus/dbus_thread_manager.h"
+#include "components/invalidation/fake_invalidation_handler.h"
 #include "components/invalidation/invalidation_service.h"
 #include "components/invalidation/invalidator_state.h"
 #include "components/invalidation/profile_invalidation_provider.h"
@@ -30,12 +31,8 @@
 #include "content/public/browser/notification_details.h"
 #include "content/public/browser/notification_service.h"
 #include "content/public/test/test_browser_thread_bundle.h"
-#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
-using testing::Mock;
-using testing::StrictMock;
-
 namespace policy {
 
 namespace {
@@ -56,16 +53,29 @@ KeyedService* BuildProfileInvalidationProvider(
 
 }  // namespace
 
-class MockConsumer : public AffiliatedInvalidationServiceProvider::Consumer {
+// A simple AffiliatedInvalidationServiceProvider::Consumer that registers a
+// syncer::FakeInvalidationHandler with the invalidation::InvalidationService
+// that is currently being made available.
+class FakeConsumer : public AffiliatedInvalidationServiceProvider::Consumer {
  public:
-  MockConsumer();
-  ~MockConsumer() override;
+  explicit FakeConsumer(AffiliatedInvalidationServiceProviderImpl* provider);
+  ~FakeConsumer() override;
+
+  // AffiliatedInvalidationServiceProvider::Consumer:
+  void OnInvalidationServiceSet(
+      invalidation::InvalidationService* invalidation_service) override;
 
-  MOCK_METHOD1(OnInvalidationServiceSet,
-               void(invalidation::InvalidationService*));
+  int GetAndClearInvalidationServiceSetCount();
+  const invalidation::InvalidationService* GetInvalidationService() const;
 
  private:
-  DISALLOW_COPY_AND_ASSIGN(MockConsumer);
+  AffiliatedInvalidationServiceProviderImpl* provider_;
+  syncer::FakeInvalidationHandler invalidation_handler_;
+
+  int invalidation_service_set_count_;

[pneubeck (no reviews), 2015/02/05 12:39:05]

you can use static initializers here:
 int invalidation_service_set_count_ = 0;
 invalidation::InvalidationService* invalidation_service_ = nullptr;

[bartfab (slow), 2015/02/05 15:15:53]

Phew, fancy C++11 :).

Done.

By the way - this would be an "initializer," not a "static initializer."

+  invalidation::InvalidationService* invalidation_service_;
+
+  DISALLOW_COPY_AND_ASSIGN(FakeConsumer);
 };
 
 class AffiliatedInvalidationServiceProviderImplTest : public testing::Test {
@@ -106,7 +116,7 @@ class AffiliatedInvalidationServiceProviderImplTest : public testing::Test {
 
  protected:
   scoped_ptr<AffiliatedInvalidationServiceProviderImpl> provider_;
-  StrictMock<MockConsumer> consumer_;
+  scoped_ptr<FakeConsumer> consumer_;
   invalidation::TiclInvalidationService* device_invalidation_service_;
   invalidation::FakeInvalidationService* profile_invalidation_service_;
 
@@ -121,10 +131,47 @@ class AffiliatedInvalidationServiceProviderImplTest : public testing::Test {
   TestingProfileManager profile_manager_;
 };
 
-MockConsumer::MockConsumer() {
+FakeConsumer::FakeConsumer(AffiliatedInvalidationServiceProviderImpl* provider)
+    : provider_(provider),
+      invalidation_service_set_count_(0),
+      invalidation_service_(nullptr) {
+  provider_->RegisterConsumer(this);
 }
 
-MockConsumer::~MockConsumer() {
+FakeConsumer::~FakeConsumer() {
+  if (invalidation_service_) {
+    invalidation_service_->UnregisterInvalidationHandler(
+        &invalidation_handler_);
+  }
+  provider_->UnregisterConsumer(this);
+
+  EXPECT_EQ(0, invalidation_service_set_count_);
+}
+
+void FakeConsumer::OnInvalidationServiceSet(
+    invalidation::InvalidationService* invalidation_service) {
+  ++invalidation_service_set_count_;
+
+  if (invalidation_service_) {
+    invalidation_service_->UnregisterInvalidationHandler(
+        &invalidation_handler_);
+  }
+
+  invalidation_service_ = invalidation_service;
+
+  if (invalidation_service_)
+    invalidation_service_->RegisterInvalidationHandler(&invalidation_handler_);
+}
+
+int FakeConsumer::GetAndClearInvalidationServiceSetCount() {
+  const int invalidation_service_set_count = invalidation_service_set_count_;
+  invalidation_service_set_count_ = 0;
+  return invalidation_service_set_count;
+}
+
+const invalidation::InvalidationService*
+FakeConsumer::GetInvalidationService() const {
+  return invalidation_service_;
 }
 
 AffiliatedInvalidationServiceProviderImplTest::
@@ -157,6 +204,7 @@ void AffiliatedInvalidationServiceProviderImplTest::SetUp() {
 }
 
 void AffiliatedInvalidationServiceProviderImplTest::TearDown() {
+  consumer_.reset();
   provider_->Shutdown();
   provider_.reset();
 
@@ -180,12 +228,9 @@ Profile* AffiliatedInvalidationServiceProviderImplTest::LogInAndReturnProfile(
 
 void AffiliatedInvalidationServiceProviderImplTest::
     LogInAsAffiliatedUserAndConnectInvalidationService() {
-  Mock::VerifyAndClearExpectations(&consumer_);
-
   // Log in as an affiliated user.
   Profile* profile = LogInAndReturnProfile(kAffiliatedUserID1);
   EXPECT_TRUE(profile);
-  Mock::VerifyAndClearExpectations(&consumer_);
 
   // Verify that a per-profile invalidation service has been created.
   profile_invalidation_service_ =
@@ -197,22 +242,18 @@ void AffiliatedInvalidationServiceProviderImplTest::
 
   // Indicate that the per-profile invalidation service has connected. Verify
   // that the consumer is informed about this.
-  EXPECT_CALL(consumer_,
-              OnInvalidationServiceSet(profile_invalidation_service_)).Times(1);
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
   profile_invalidation_service_->SetInvalidatorState(
       syncer::INVALIDATIONS_ENABLED);
-  Mock::VerifyAndClearExpectations(&consumer_);
+  EXPECT_EQ(1, consumer_->GetAndClearInvalidationServiceSetCount());
+  EXPECT_EQ(profile_invalidation_service_, consumer_->GetInvalidationService());
 
   // Verify that the device-global invalidation service has been destroyed.
   EXPECT_FALSE(provider_->GetDeviceInvalidationServiceForTest());
-
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 void AffiliatedInvalidationServiceProviderImplTest::
     LogInAsUnaffiliatedUserAndConnectInvalidationService() {
-  Mock::VerifyAndClearExpectations(&consumer_);
-
   // Log in as an unaffiliated user.
   Profile* profile = LogInAndReturnProfile(kUnaffiliatedUserID);
   EXPECT_TRUE(profile);
@@ -229,17 +270,14 @@ void AffiliatedInvalidationServiceProviderImplTest::
   // that the consumer is not called back.
   profile_invalidation_service_->SetInvalidatorState(
       syncer::INVALIDATIONS_ENABLED);
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Verify that the device-global invalidation service still exists.
   EXPECT_TRUE(provider_->GetDeviceInvalidationServiceForTest());
-
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 void AffiliatedInvalidationServiceProviderImplTest::
     ConnectDeviceGlobalInvalidationService() {
-  Mock::VerifyAndClearExpectations(&consumer_);
-
   // Verify that a device-global invalidation service has been created.
   device_invalidation_service_ =
       provider_->GetDeviceInvalidationServiceForTest();
@@ -247,30 +285,27 @@ void AffiliatedInvalidationServiceProviderImplTest::
 
   // Indicate that the device-global invalidation service has connected. Verify
   // that the consumer is informed about this.
-  EXPECT_CALL(consumer_, OnInvalidationServiceSet(device_invalidation_service_))
-      .Times(1);
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
   device_invalidation_service_->OnInvalidatorStateChange(
       syncer::INVALIDATIONS_ENABLED);
-
-  Mock::VerifyAndClearExpectations(&consumer_);
+  EXPECT_EQ(1, consumer_->GetAndClearInvalidationServiceSetCount());
+  EXPECT_EQ(device_invalidation_service_, consumer_->GetInvalidationService());
 }
 
 void AffiliatedInvalidationServiceProviderImplTest::
     DisconnectPerProfileInvalidationService() {
-  Mock::VerifyAndClearExpectations(&consumer_);
-
   ASSERT_TRUE(profile_invalidation_service_);
 
   // Indicate that the per-profile invalidation service has disconnected. Verify
   // that the consumer is informed about this.
-  EXPECT_CALL(consumer_, OnInvalidationServiceSet(nullptr)).Times(1);
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
   profile_invalidation_service_->SetInvalidatorState(
       syncer::INVALIDATION_CREDENTIALS_REJECTED);
+  EXPECT_EQ(1, consumer_->GetAndClearInvalidationServiceSetCount());
+  EXPECT_EQ(nullptr, consumer_->GetInvalidationService());
 
   // Verify that a device-global invalidation service has been created.
   EXPECT_TRUE(provider_->GetDeviceInvalidationServiceForTest());
-
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 invalidation::FakeInvalidationService*
@@ -310,7 +345,8 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
        UseDeviceInvalidationService) {
   // Register a consumer. Verify that the consumer is not called back
   // immediately as no connected invalidation service exists yet.
-  provider_->RegisterConsumer(&consumer_);
+  consumer_.reset(new FakeConsumer(provider_.get()));
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Indicate that the device-global invalidation service connected. Verify that
   // that the consumer is informed about this.
@@ -318,17 +354,14 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
 
   // Indicate that the device-global invalidation service has disconnected.
   // Verify that the consumer is informed about this.
-  EXPECT_CALL(consumer_, OnInvalidationServiceSet(nullptr)).Times(1);
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
   device_invalidation_service_->OnInvalidatorStateChange(
       syncer::INVALIDATION_CREDENTIALS_REJECTED);
-  Mock::VerifyAndClearExpectations(&consumer_);
+  EXPECT_EQ(1, consumer_->GetAndClearInvalidationServiceSetCount());
+  EXPECT_EQ(nullptr, consumer_->GetInvalidationService());
 
   // Verify that the device-global invalidation service still exists.
   EXPECT_TRUE(provider_->GetDeviceInvalidationServiceForTest());
-
-  // Unregister the consumer.
-  provider_->UnregisterConsumer(&consumer_);
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 // A consumer is registered with the AffiliatedInvalidationServiceProviderImpl.
@@ -338,7 +371,8 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
        UseAffiliatedProfileInvalidationService) {
   // Register a consumer. Verify that the consumer is not called back

[pneubeck (no reviews), 2015/02/05 12:39:05]

I know you didn't add it in this CL, but instead of refactoring the same code X
times is you could just put it into a separate tiny test.

TEST_F(AffiliatedInvalidationServiceProviderImplTest,
NoNotifcationIfNoServiceConnected) {
  // Register a consumer. Verify that the consumer is not called back
  // immediately as no connected invalidation service exists yet.
  consumer_.reset(new FakeConsumer(provider_.get()));
  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
}

The initial comment and
  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
could be removed from all tests.
consumer_.reset(...) is self explanatory.
  

[bartfab (slow), 2015/02/05 15:15:53]

Done.

   // immediately as no connected invalidation service exists yet.
-  provider_->RegisterConsumer(&consumer_);
+  consumer_.reset(new FakeConsumer(provider_.get()));
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Verify that a device-global invalidation service has been created.
   EXPECT_TRUE(provider_->GetDeviceInvalidationServiceForTest());
@@ -353,10 +387,6 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
   // disconnected. Verify that the consumer is informed about this and a
   // device-global invalidation service is created.
   DisconnectPerProfileInvalidationService();
-
-  // Unregister the consumer.
-  provider_->UnregisterConsumer(&consumer_);
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 // A consumer is registered with the AffiliatedInvalidationServiceProviderImpl.
@@ -366,7 +396,8 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
        DoNotUseUnaffiliatedProfileInvalidationService) {
   // Register a consumer. Verify that the consumer is not called back
   // immediately as no connected invalidation service exists yet.
-  provider_->RegisterConsumer(&consumer_);
+  consumer_.reset(new FakeConsumer(provider_.get()));
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Verify that a device-global invalidation service has been created.
   EXPECT_TRUE(provider_->GetDeviceInvalidationServiceForTest());
@@ -376,10 +407,6 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
   // service is ignored and the device-global invalidation service is not
   // destroyed.
   LogInAsUnaffiliatedUserAndConnectInvalidationService();
-
-  // Unregister the consumer.
-  provider_->UnregisterConsumer(&consumer_);
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 // A consumer is registered with the AffiliatedInvalidationServiceProviderImpl.
@@ -391,7 +418,8 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
        SwitchToAffiliatedProfileInvalidationService) {
   // Register a consumer. Verify that the consumer is not called back
   // immediately as no connected invalidation service exists yet.
-  provider_->RegisterConsumer(&consumer_);
+  consumer_.reset(new FakeConsumer(provider_.get()));
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Indicate that the device-global invalidation service connected. Verify that
   // that the consumer is informed about this.
@@ -402,10 +430,6 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
   // made available to the |consumer_| and the device-global invalidation
   // service is destroyed.
   LogInAsAffiliatedUserAndConnectInvalidationService();
-
-  // Unregister the consumer.
-  provider_->UnregisterConsumer(&consumer_);
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 // A consumer is registered with the AffiliatedInvalidationServiceProviderImpl.
@@ -418,7 +442,8 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
        DoNotSwitchToUnaffiliatedProfileInvalidationService) {
   // Register a consumer. Verify that the consumer is not called back
   // immediately as no connected invalidation service exists yet.
-  provider_->RegisterConsumer(&consumer_);
+  consumer_.reset(new FakeConsumer(provider_.get()));
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Indicate that the device-global invalidation service connected. Verify that
   // that the consumer is informed about this.
@@ -429,10 +454,6 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
   // service is ignored and the device-global invalidation service is not
   // destroyed.
   LogInAsUnaffiliatedUserAndConnectInvalidationService();
-
-  // Unregister the consumer.
-  provider_->UnregisterConsumer(&consumer_);
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 // A consumer is registered with the AffiliatedInvalidationServiceProviderImpl.
@@ -445,7 +466,8 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
        SwitchToDeviceInvalidationService) {
   // Register a consumer. Verify that the consumer is not called back
   // immediately as no connected invalidation service exists yet.
-  provider_->RegisterConsumer(&consumer_);
+  consumer_.reset(new FakeConsumer(provider_.get()));
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Verify that a device-global invalidation service has been created.
   EXPECT_TRUE(provider_->GetDeviceInvalidationServiceForTest());
@@ -464,10 +486,6 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
   // Indicate that the device-global invalidation service connected. Verify that
   // that the consumer is informed about this.
   ConnectDeviceGlobalInvalidationService();
-
-  // Unregister the consumer.
-  provider_->UnregisterConsumer(&consumer_);
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 // A consumer is registered with the AffiliatedInvalidationServiceProviderImpl.
@@ -481,7 +499,8 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
        SwitchBetweenAffiliatedProfileInvalidationServices) {
   // Register a consumer. Verify that the consumer is not called back
   // immediately as no connected invalidation service exists yet.
-  provider_->RegisterConsumer(&consumer_);
+  consumer_.reset(new FakeConsumer(provider_.get()));
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Verify that a device-global invalidation service has been created.
   EXPECT_TRUE(provider_->GetDeviceInvalidationServiceForTest());
@@ -509,25 +528,21 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
   // connected. Verify that the consumer is not called back.
   second_profile_invalidation_service->SetInvalidatorState(
       syncer::INVALIDATIONS_ENABLED);
-  Mock::VerifyAndClearExpectations(&consumer_);
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Indicate that the first user's per-profile invalidation service has
   // disconnected. Verify that the consumer is informed that the second user's
   // per-profile invalidation service should be used instead of the first
   // user's.
-  EXPECT_CALL(consumer_,
-              OnInvalidationServiceSet(second_profile_invalidation_service))
-      .Times(1);
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
   profile_invalidation_service_->SetInvalidatorState(
       syncer::INVALIDATION_CREDENTIALS_REJECTED);
-  Mock::VerifyAndClearExpectations(&consumer_);
+  EXPECT_EQ(1, consumer_->GetAndClearInvalidationServiceSetCount());
+  EXPECT_EQ(second_profile_invalidation_service,
+            consumer_->GetInvalidationService());
 
   // Verify that the device-global invalidation service still does not exist.
   EXPECT_FALSE(provider_->GetDeviceInvalidationServiceForTest());
-
-  // Unregister the consumer.
-  provider_->UnregisterConsumer(&consumer_);
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 // A consumer is registered with the AffiliatedInvalidationServiceProviderImpl.
@@ -541,7 +556,8 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest,
 TEST_F(AffiliatedInvalidationServiceProviderImplTest, MultipleConsumers) {
   // Register a first consumer. Verify that the consumer is not called back
   // immediately as no connected invalidation service exists yet.
-  provider_->RegisterConsumer(&consumer_);
+  consumer_.reset(new FakeConsumer(provider_.get()));
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Indicate that the device-global invalidation service connected. Verify that
   // that the consumer is informed about this.
@@ -549,25 +565,22 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest, MultipleConsumers) {
 
   // Register a second consumer. Verify that the consumer is called back
   // immediately as a connected invalidation service is available.
-  StrictMock<MockConsumer> second_consumer;
-  EXPECT_CALL(second_consumer,
-              OnInvalidationServiceSet(device_invalidation_service_)).Times(1);
-  provider_->RegisterConsumer(&second_consumer);
-  Mock::VerifyAndClearExpectations(&second_consumer);
+  scoped_ptr<FakeConsumer> second_consumer(new FakeConsumer(provider_.get()));
+  EXPECT_EQ(1, second_consumer->GetAndClearInvalidationServiceSetCount());
+  EXPECT_EQ(device_invalidation_service_,
+            second_consumer->GetInvalidationService());
 
   // Unregister the first consumer.
-  provider_->UnregisterConsumer(&consumer_);
+  consumer_.reset();
 
   // Verify that the device-global invalidation service still exists.
   EXPECT_TRUE(provider_->GetDeviceInvalidationServiceForTest());
 
   // Unregister the second consumer.
-  provider_->UnregisterConsumer(&second_consumer);
+  second_consumer.reset();
 
   // Verify that the device-global invalidation service has been destroyed.
   EXPECT_FALSE(provider_->GetDeviceInvalidationServiceForTest());
-  Mock::VerifyAndClearExpectations(&consumer_);
-  Mock::VerifyAndClearExpectations(&second_consumer);
 }
 
 // A consumer is registered with the AffiliatedInvalidationServiceProviderImpl.
@@ -581,7 +594,8 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest, MultipleConsumers) {
 TEST_F(AffiliatedInvalidationServiceProviderImplTest, NoServiceAfterShutdown) {
   // Register a consumer. Verify that the consumer is not called back
   // immediately as no connected invalidation service exists yet.
-  provider_->RegisterConsumer(&consumer_);
+  consumer_.reset(new FakeConsumer(provider_.get()));
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Verify that a device-global invalidation service has been created.
   EXPECT_TRUE(provider_->GetDeviceInvalidationServiceForTest());
@@ -594,9 +608,10 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest, NoServiceAfterShutdown) {
 
   // Shut down the |provider_|. Verify that the |consumer_| is informed that no
   // invalidation service is available for use anymore.
-  EXPECT_CALL(consumer_, OnInvalidationServiceSet(nullptr)).Times(1);
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
   provider_->Shutdown();
-  Mock::VerifyAndClearExpectations(&consumer_);
+  EXPECT_EQ(1, consumer_->GetAndClearInvalidationServiceSetCount());
+  EXPECT_EQ(nullptr, consumer_->GetInvalidationService());
 
   // Verify that the device-global invalidation service still does not exist.
   EXPECT_FALSE(provider_->GetDeviceInvalidationServiceForTest());
@@ -617,13 +632,10 @@ TEST_F(AffiliatedInvalidationServiceProviderImplTest, NoServiceAfterShutdown) {
   // connected. Verify that the consumer is not called back.
   second_profile_invalidation_service->SetInvalidatorState(
       syncer::INVALIDATIONS_ENABLED);
+  EXPECT_EQ(0, consumer_->GetAndClearInvalidationServiceSetCount());
 
   // Verify that the device-global invalidation service still does not exist.
   EXPECT_FALSE(provider_->GetDeviceInvalidationServiceForTest());
-
-  // Unregister the consumer.
-  provider_->UnregisterConsumer(&consumer_);
-  Mock::VerifyAndClearExpectations(&consumer_);
 }
 
 }  // namespace policy