[GCM] Adding registration ID request and tests to GCM Account Mapper

components/gcm_driver/gcm_account_mapper_unittest.cc

 // Copyright 2014 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "components/gcm_driver/gcm_account_mapper.h"
 
 #include "base/test/simple_test_clock.h"
 #include "base/time/time.h"
 #include "components/gcm_driver/fake_gcm_driver.h"
 #include "google_apis/gcm/engine/account_mapping.h"
@@ skipping 69 matching lines @@
 
   CustomFakeGCMDriver();
   virtual ~CustomFakeGCMDriver();
 
   virtual void UpdateAccountMapping(
       const AccountMapping& account_mapping) OVERRIDE;
   virtual void RemoveAccountMapping(const std::string& account_id) OVERRIDE;
   virtual void AddAppHandler(const std::string& app_id,
                              GCMAppHandler* handler) OVERRIDE;
   virtual void RemoveAppHandler(const std::string& app_id) OVERRIDE;
+  virtual void RegisterImpl(
+      const std::string& app_id,
+      const std::vector<std::string>& sender_ids) OVERRIDE;
 
+  void CompleteRegister(const std::string& registration_id,
+                        GCMClient::Result result);
   void CompleteSend(const std::string& message_id, GCMClient::Result result);
   void AcknowledgeSend(const std::string& message_id);
   void MessageSendError(const std::string& message_id);
 
   void CompleteSendAllMessages();
   void AcknowledgeSendAllMessages();
 
   void SetLastMessageAction(const std::string& message_id,
                             LastMessageAction action);
   void Clear();
 
   const AccountMapping& last_account_mapping() const {
     return account_mapping_;
   }
   const std::string& last_message_id() const { return last_message_id_; }
   const std::string& last_removed_account_id() const {
     return last_removed_account_id_;
   }
   LastMessageAction last_action() const { return last_action_; }
+  bool registration_id_requested() const { return registration_id_requested_; }
 
  protected:
   virtual void SendImpl(const std::string& app_id,
                         const std::string& receiver_id,
                         const GCMClient::OutgoingMessage& message) OVERRIDE;
 
  private:
   AccountMapping account_mapping_;
   std::string last_message_id_;
   std::string last_removed_account_id_;
   LastMessageAction last_action_;
   std::map<std::string, LastMessageAction> all_messages_;
+  bool registration_id_requested_;
 };
 
-CustomFakeGCMDriver::CustomFakeGCMDriver() : last_action_(NONE) {
+CustomFakeGCMDriver::CustomFakeGCMDriver()
+    : last_action_(NONE), registration_id_requested_(false) {
 }
 
 CustomFakeGCMDriver::~CustomFakeGCMDriver() {
 }
 
 void CustomFakeGCMDriver::UpdateAccountMapping(
     const AccountMapping& account_mapping) {
   account_mapping_.email = account_mapping.email;
   account_mapping_.account_id = account_mapping.account_id;
   account_mapping_.access_token = account_mapping.access_token;
   account_mapping_.status = account_mapping.status;
   account_mapping_.status_change_timestamp =
       account_mapping.status_change_timestamp;
   account_mapping_.last_message_id = account_mapping.last_message_id;
 }
 
 void CustomFakeGCMDriver::RemoveAccountMapping(const std::string& account_id) {
   last_removed_account_id_ = account_id;
 }
 
 void CustomFakeGCMDriver::AddAppHandler(const std::string& app_id,
                                         GCMAppHandler* handler) {
   GCMDriver::AddAppHandler(app_id, handler);
 }
 
 void CustomFakeGCMDriver::RemoveAppHandler(const std::string& app_id) {
   GCMDriver::RemoveAppHandler(app_id);
 }
 
+void CustomFakeGCMDriver::RegisterImpl(
+    const std::string& app_id,
+    const std::vector<std::string>& sender_ids) {
+  DCHECK_EQ(kGCMAccountMapperAppId, app_id);
+  DCHECK_EQ(1u, sender_ids.size());
+  DCHECK_EQ(kGCMAccountMapperSenderId, sender_ids[0]);
+  registration_id_requested_ = true;
+}
+
+void CustomFakeGCMDriver::CompleteRegister(const std::string& registration_id,
+                                           GCMClient::Result result) {
+  RegisterFinished(kGCMAccountMapperAppId, registration_id, result);
+}
+
 void CustomFakeGCMDriver::CompleteSend(const std::string& message_id,
                                        GCMClient::Result result) {
   SendFinished(kGCMAccountMapperAppId, message_id, result);
   SetLastMessageAction(message_id, SEND_FINISHED);
 }
 
 void CustomFakeGCMDriver::AcknowledgeSend(const std::string& message_id) {
   GetAppHandler(kGCMAccountMapperAppId)
       ->OnSendAcknowledged(kGCMAccountMapperAppId, message_id);
   SetLastMessageAction(message_id, SEND_ACKNOWLEDGED);
@@ skipping 34 matching lines @@
     if (iter->second == SEND_FINISHED)
       AcknowledgeSend(iter->first);
   }
 }
 
 void CustomFakeGCMDriver::Clear() {
   account_mapping_ = AccountMapping();
   last_message_id_.clear();
   last_removed_account_id_.clear();
   last_action_ = NONE;
+  registration_id_requested_ = false;
 }
 
 void CustomFakeGCMDriver::SetLastMessageAction(const std::string& message_id,
                                                LastMessageAction action) {
   last_action_ = action;
   last_message_id_ = message_id;
   all_messages_[message_id] = action;
 }
 
 }  // namespace
@@ skipping 31 matching lines @@
 void GCMAccountMapperTest::Restart() {
   if (account_mapper_)
     account_mapper_->ShutdownHandler();
   account_mapper_.reset(new GCMAccountMapper(&gcm_driver_));
   scoped_ptr<base::SimpleTestClock> clock(new base::SimpleTestClock);
   clock_ = clock.get();
   account_mapper_->SetClockForTesting(clock.PassAs<base::Clock>());
 }
 
 // Tests the initialization of account mappings (from the store) when empty.
+// It also checks that initialization triggers registration ID request.
 TEST_F(GCMAccountMapperTest, InitializeAccountMappingsEmpty) {
-  GCMAccountMapper::AccountMappings account_mappings;
-  mapper()->Initialize(account_mappings, "");
+  EXPECT_FALSE(gcm_driver().registration_id_requested());
+  mapper()->Initialize(GCMAccountMapper::AccountMappings());
   EXPECT_TRUE(GetAccounts().empty());
+  EXPECT_TRUE(gcm_driver().registration_id_requested());
+}
+
+// Tests that registration is retried, unless there is async operation pending.
+TEST_F(GCMAccountMapperTest, RegistrationRetryUponFailure) {
+  mapper()->Initialize(GCMAccountMapper::AccountMappings());
+  EXPECT_TRUE(gcm_driver().registration_id_requested());
+  gcm_driver().Clear();
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::UNKNOWN_ERROR);
+  EXPECT_TRUE(gcm_driver().registration_id_requested());
+  gcm_driver().Clear();
+  gcm_driver().CompleteRegister(kRegistrationId,
+                                GCMClient::ASYNC_OPERATION_PENDING);
+  EXPECT_FALSE(gcm_driver().registration_id_requested());
 }
 
 // Tests the initialization of account mappings (from the store).
 TEST_F(GCMAccountMapperTest, InitializeAccountMappings) {
   GCMAccountMapper::AccountMappings account_mappings;
   AccountMapping account_mapping1 = MakeAccountMapping("acc_id1",
                                                        AccountMapping::MAPPED,
                                                        base::Time::Now(),
                                                        std::string());
   AccountMapping account_mapping2 = MakeAccountMapping("acc_id2",
                                                        AccountMapping::ADDING,
                                                        base::Time::Now(),
                                                        "add_message_1");
   account_mappings.push_back(account_mapping1);
   account_mappings.push_back(account_mapping2);
 
-  mapper()->Initialize(account_mappings, "");
+  mapper()->Initialize(account_mappings);
 
   GCMAccountMapper::AccountMappings mappings = GetAccounts();
   EXPECT_EQ(2UL, mappings.size());
   GCMAccountMapper::AccountMappings::const_iterator iter = mappings.begin();
 
   EXPECT_EQ(account_mapping1.account_id, iter->account_id);
   EXPECT_EQ(account_mapping1.email, iter->email);
   EXPECT_TRUE(account_mapping1.access_token.empty());
   EXPECT_EQ(account_mapping1.status, iter->status);
   EXPECT_EQ(account_mapping1.status_change_timestamp,
             iter->status_change_timestamp);
   EXPECT_TRUE(account_mapping1.last_message_id.empty());
 
   ++iter;
   EXPECT_EQ(account_mapping2.account_id, iter->account_id);
   EXPECT_EQ(account_mapping2.email, iter->email);
   EXPECT_TRUE(account_mapping2.access_token.empty());
   EXPECT_EQ(account_mapping2.status, iter->status);
   EXPECT_EQ(account_mapping2.status_change_timestamp,
             iter->status_change_timestamp);
   EXPECT_EQ(account_mapping2.last_message_id, iter->last_message_id);
 }
 
+// Tests that account tokens are not processed until registration ID is
+// available.
+TEST_F(GCMAccountMapperTest, SetAccountTokensOnlyWorksWithRegisterationId) {
+  // Start with empty list.
+  mapper()->Initialize(GCMAccountMapper::AccountMappings());
+
+  std::vector<GCMClient::AccountTokenInfo> account_tokens;
+  account_tokens.push_back(MakeAccountTokenInfo("acc_id"));
+  mapper()->SetAccountTokens(account_tokens);
+
+  EXPECT_TRUE(GetAccounts().empty());
+
+  account_tokens.clear();
+  account_tokens.push_back(MakeAccountTokenInfo("acc_id1"));
+  account_tokens.push_back(MakeAccountTokenInfo("acc_id2"));
+  mapper()->SetAccountTokens(account_tokens);
+
+  EXPECT_TRUE(GetAccounts().empty());
+
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
+
+  GCMAccountMapper::AccountMappings mappings = GetAccounts();
+  EXPECT_EQ(2UL, mappings.size());
+  EXPECT_EQ("acc_id1", mappings[0].account_id);
+  EXPECT_EQ("acc_id2", mappings[1].account_id);
+}
+
 // Tests the part where a new account is added with a token, to the point when
 // GCM message is sent.
 TEST_F(GCMAccountMapperTest, AddMappingToMessageSent) {
-  mapper()->Initialize(GCMAccountMapper::AccountMappings(), kRegistrationId);
+  mapper()->Initialize(GCMAccountMapper::AccountMappings());
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
 
   std::vector<GCMClient::AccountTokenInfo> account_tokens;
   GCMClient::AccountTokenInfo account_token = MakeAccountTokenInfo("acc_id");
   account_tokens.push_back(account_token);
   mapper()->SetAccountTokens(account_tokens);
 
   GCMAccountMapper::AccountMappings mappings = GetAccounts();
   EXPECT_EQ(1UL, mappings.size());
   GCMAccountMapper::AccountMappings::const_iterator iter = mappings.begin();
   EXPECT_EQ("acc_id", iter->account_id);
   EXPECT_EQ("acc_id@gmail.com", iter->email);
   EXPECT_EQ("acc_id_token", iter->access_token);
   EXPECT_EQ(AccountMapping::NEW, iter->status);
   EXPECT_EQ(base::Time(), iter->status_change_timestamp);
 
   EXPECT_TRUE(!gcm_driver().last_message_id().empty());
 }
 
 // Tests the part where GCM message is successfully queued.
 TEST_F(GCMAccountMapperTest, AddMappingMessageQueued) {
-  mapper()->Initialize(GCMAccountMapper::AccountMappings(), kRegistrationId);
+  mapper()->Initialize(GCMAccountMapper::AccountMappings());
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
 
   std::vector<GCMClient::AccountTokenInfo> account_tokens;
   GCMClient::AccountTokenInfo account_token = MakeAccountTokenInfo("acc_id");
   account_tokens.push_back(account_token);
   mapper()->SetAccountTokens(account_tokens);
 
   clock()->SetNow(base::Time::Now());
   gcm_driver().CompleteSend(gcm_driver().last_message_id(), GCMClient::SUCCESS);
 
   EXPECT_EQ(account_token.email, gcm_driver().last_account_mapping().email);
@@ skipping 12 matching lines @@
   EXPECT_EQ(account_token.email, iter->email);
   EXPECT_EQ(account_token.account_id, iter->account_id);
   EXPECT_EQ(account_token.access_token, iter->access_token);
   EXPECT_EQ(AccountMapping::ADDING, iter->status);
   EXPECT_EQ(clock()->Now(), iter->status_change_timestamp);
   EXPECT_EQ(gcm_driver().last_message_id(), iter->last_message_id);
 }
 
 // Tests status change from ADDING to MAPPED (Message is acknowledged).
 TEST_F(GCMAccountMapperTest, AddMappingMessageAcknowledged) {
-  mapper()->Initialize(GCMAccountMapper::AccountMappings(), kRegistrationId);
+  mapper()->Initialize(GCMAccountMapper::AccountMappings());
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
 
   std::vector<GCMClient::AccountTokenInfo> account_tokens;
   GCMClient::AccountTokenInfo account_token = MakeAccountTokenInfo("acc_id");
   account_tokens.push_back(account_token);
   mapper()->SetAccountTokens(account_tokens);
 
   clock()->SetNow(base::Time::Now());
   gcm_driver().CompleteSend(gcm_driver().last_message_id(), GCMClient::SUCCESS);
   clock()->SetNow(base::Time::Now());
   gcm_driver().AcknowledgeSend(gcm_driver().last_message_id());
@@ skipping 14 matching lines @@
   EXPECT_EQ(account_token.account_id, iter->account_id);
   EXPECT_EQ(account_token.access_token, iter->access_token);
   EXPECT_EQ(AccountMapping::MAPPED, iter->status);
   EXPECT_EQ(clock()->Now(), iter->status_change_timestamp);
   EXPECT_TRUE(iter->last_message_id.empty());
 }
 
 // Tests status change form ADDING to MAPPED (When message was acknowledged,
 // after Chrome was restarted).
 TEST_F(GCMAccountMapperTest, AddMappingMessageAckedAfterRestart) {
-  mapper()->Initialize(GCMAccountMapper::AccountMappings(), kRegistrationId);
+  mapper()->Initialize(GCMAccountMapper::AccountMappings());
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
 
   std::vector<GCMClient::AccountTokenInfo> account_tokens;
   GCMClient::AccountTokenInfo account_token = MakeAccountTokenInfo("acc_id");
   account_tokens.push_back(account_token);
   mapper()->SetAccountTokens(account_tokens);
 
   clock()->SetNow(base::Time::Now());
   gcm_driver().CompleteSend(gcm_driver().last_message_id(), GCMClient::SUCCESS);
 
   Restart();
   GCMAccountMapper::AccountMappings stored_mappings;
   stored_mappings.push_back(gcm_driver().last_account_mapping());
-  mapper()->Initialize(stored_mappings, kRegistrationId);
+  mapper()->Initialize(stored_mappings);
 
   clock()->SetNow(base::Time::Now());
   gcm_driver().AcknowledgeSend(gcm_driver().last_message_id());
 
   EXPECT_EQ(account_token.email, gcm_driver().last_account_mapping().email);
   EXPECT_EQ(account_token.account_id,
             gcm_driver().last_account_mapping().account_id);
   EXPECT_EQ(account_token.access_token,
             gcm_driver().last_account_mapping().access_token);
   EXPECT_EQ(AccountMapping::MAPPED, gcm_driver().last_account_mapping().status);
   EXPECT_EQ(clock()->Now(),
             gcm_driver().last_account_mapping().status_change_timestamp);
   EXPECT_TRUE(gcm_driver().last_account_mapping().last_message_id.empty());
 
   GCMAccountMapper::AccountMappings mappings = GetAccounts();
   GCMAccountMapper::AccountMappings::const_iterator iter = mappings.begin();
   EXPECT_EQ(account_token.email, iter->email);
   EXPECT_EQ(account_token.account_id, iter->account_id);
   EXPECT_EQ(account_token.access_token, iter->access_token);
   EXPECT_EQ(AccountMapping::MAPPED, iter->status);
   EXPECT_EQ(clock()->Now(), iter->status_change_timestamp);
   EXPECT_TRUE(iter->last_message_id.empty());
 }
 
 // Tests a case when ADD message times out for a new account.
 TEST_F(GCMAccountMapperTest, AddMappingMessageSendErrorForNewAccount) {
-  mapper()->Initialize(GCMAccountMapper::AccountMappings(), kRegistrationId);
+  mapper()->Initialize(GCMAccountMapper::AccountMappings());
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
 
   std::vector<GCMClient::AccountTokenInfo> account_tokens;
   GCMClient::AccountTokenInfo account_token = MakeAccountTokenInfo("acc_id");
   account_tokens.push_back(account_token);
   mapper()->SetAccountTokens(account_tokens);
 
   clock()->SetNow(base::Time::Now());
   gcm_driver().CompleteSend(gcm_driver().last_message_id(), GCMClient::SUCCESS);
 
   clock()->SetNow(base::Time::Now());
@@ skipping 11 matching lines @@
 TEST_F(GCMAccountMapperTest, AddMappingMessageSendErrorForMappedAccount) {
   // Start with one account that is mapped.
   base::Time status_change_timestamp = base::Time::Now();
   AccountMapping mapping = MakeAccountMapping("acc_id",
                                               AccountMapping::MAPPED,
                                               status_change_timestamp,
                                               "add_message_id");
 
   GCMAccountMapper::AccountMappings stored_mappings;
   stored_mappings.push_back(mapping);
-  mapper()->Initialize(stored_mappings, kRegistrationId);
+  mapper()->Initialize(stored_mappings);
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
 
   clock()->SetNow(base::Time::Now());
   gcm_driver().MessageSendError("add_message_id");
 
   // No new message is sent because of the send error, as the token is stale.
   // Because the account was new, the entry should be deleted.
   EXPECT_TRUE(gcm_driver().last_message_id().empty());
 
   GCMAccountMapper::AccountMappings mappings = GetAccounts();
   GCMAccountMapper::AccountMappings::const_iterator iter = mappings.begin();
   EXPECT_EQ(mapping.email, iter->email);
   EXPECT_EQ(mapping.account_id, iter->account_id);
   EXPECT_EQ(mapping.access_token, iter->access_token);
   EXPECT_EQ(AccountMapping::MAPPED, iter->status);
   EXPECT_EQ(status_change_timestamp, iter->status_change_timestamp);
   EXPECT_TRUE(iter->last_message_id.empty());
 }
 
 // Tests that a missing token for an account will trigger removing of that
 // account. This test goes only until the message is passed to GCM.
 TEST_F(GCMAccountMapperTest, RemoveMappingToMessageSent) {
   // Start with one account that is mapped.
   AccountMapping mapping = MakeAccountMapping("acc_id",
                                               AccountMapping::MAPPED,
                                               base::Time::Now(),
                                               std::string());
 
   GCMAccountMapper::AccountMappings stored_mappings;
   stored_mappings.push_back(mapping);
-  mapper()->Initialize(stored_mappings, kRegistrationId);
+  mapper()->Initialize(stored_mappings);
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
   clock()->SetNow(base::Time::Now());
 
   mapper()->SetAccountTokens(std::vector<GCMClient::AccountTokenInfo>());
 
   EXPECT_EQ(mapping.account_id, gcm_driver().last_account_mapping().account_id);
   EXPECT_EQ(mapping.email, gcm_driver().last_account_mapping().email);
   EXPECT_EQ(AccountMapping::REMOVING,
             gcm_driver().last_account_mapping().status);
   EXPECT_EQ(clock()->Now(),
             gcm_driver().last_account_mapping().status_change_timestamp);
@@ skipping 13 matching lines @@
 // account. This test goes until the message is queued by GCM.
 TEST_F(GCMAccountMapperTest, RemoveMappingMessageQueued) {
   // Start with one account that is mapped.
   AccountMapping mapping = MakeAccountMapping("acc_id",
                                               AccountMapping::MAPPED,
                                               base::Time::Now(),
                                               std::string());
 
   GCMAccountMapper::AccountMappings stored_mappings;
   stored_mappings.push_back(mapping);
-  mapper()->Initialize(stored_mappings, kRegistrationId);
+  mapper()->Initialize(stored_mappings);
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
   clock()->SetNow(base::Time::Now());
   base::Time status_change_timestamp = clock()->Now();
 
   mapper()->SetAccountTokens(std::vector<GCMClient::AccountTokenInfo>());
   clock()->SetNow(base::Time::Now());
   gcm_driver().CompleteSend(gcm_driver().last_message_id(), GCMClient::SUCCESS);
 
   EXPECT_EQ(mapping.account_id, gcm_driver().last_account_mapping().account_id);
   EXPECT_EQ(mapping.email, gcm_driver().last_account_mapping().email);
   EXPECT_EQ(AccountMapping::REMOVING,
@@ skipping 19 matching lines @@
 // account mapping being completely gone.
 TEST_F(GCMAccountMapperTest, RemoveMappingMessageAcknowledged) {
   // Start with one account that is mapped.
   AccountMapping mapping = MakeAccountMapping("acc_id",
                                               AccountMapping::MAPPED,
                                               base::Time::Now(),
                                               std::string());
 
   GCMAccountMapper::AccountMappings stored_mappings;
   stored_mappings.push_back(mapping);
-  mapper()->Initialize(stored_mappings, kRegistrationId);
+  mapper()->Initialize(stored_mappings);
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
   clock()->SetNow(base::Time::Now());
 
   mapper()->SetAccountTokens(std::vector<GCMClient::AccountTokenInfo>());
   gcm_driver().CompleteSend(gcm_driver().last_message_id(), GCMClient::SUCCESS);
   gcm_driver().AcknowledgeSend(gcm_driver().last_message_id());
 
   EXPECT_EQ(mapping.account_id, gcm_driver().last_removed_account_id());
 
   GCMAccountMapper::AccountMappings mappings = GetAccounts();
   EXPECT_TRUE(mappings.empty());
 }
 
 // Tests that account removing proceeds, when a removing message is acked after
 // Chrome was restarted.
 TEST_F(GCMAccountMapperTest, RemoveMappingMessageAckedAfterRestart) {
   // Start with one account that is mapped.
   AccountMapping mapping = MakeAccountMapping("acc_id",
                                               AccountMapping::REMOVING,
                                               base::Time::Now(),
                                               "remove_message_id");
 
   GCMAccountMapper::AccountMappings stored_mappings;
   stored_mappings.push_back(mapping);
-  mapper()->Initialize(stored_mappings, kRegistrationId);
+  mapper()->Initialize(stored_mappings);
 
   gcm_driver().AcknowledgeSend("remove_message_id");
 
   EXPECT_EQ(mapping.account_id, gcm_driver().last_removed_account_id());
 
   GCMAccountMapper::AccountMappings mappings = GetAccounts();
   EXPECT_TRUE(mappings.empty());
 }
 
 // Tests that account removing proceeds, when a removing message is acked after
 // Chrome was restarted.
 TEST_F(GCMAccountMapperTest, RemoveMappingMessageSendError) {
   // Start with one account that is mapped.
   base::Time status_change_timestamp = base::Time::Now();
   AccountMapping mapping = MakeAccountMapping("acc_id",
                                               AccountMapping::REMOVING,
                                               status_change_timestamp,
                                               "remove_message_id");
 
   GCMAccountMapper::AccountMappings stored_mappings;
   stored_mappings.push_back(mapping);
-  mapper()->Initialize(stored_mappings, kRegistrationId);
+  mapper()->Initialize(stored_mappings);
 
   clock()->SetNow(base::Time::Now());
   gcm_driver().MessageSendError("remove_message_id");
 
   EXPECT_TRUE(gcm_driver().last_removed_account_id().empty());
 
   EXPECT_EQ(mapping.account_id, gcm_driver().last_account_mapping().account_id);
   EXPECT_EQ(mapping.email, gcm_driver().last_account_mapping().email);
   EXPECT_EQ(AccountMapping::REMOVING,
             gcm_driver().last_account_mapping().status);
   EXPECT_EQ(status_change_timestamp,
             gcm_driver().last_account_mapping().status_change_timestamp);
   // Message is not persisted, until send is completed.
   EXPECT_TRUE(gcm_driver().last_account_mapping().last_message_id.empty());
 
   GCMAccountMapper::AccountMappings mappings = GetAccounts();
   GCMAccountMapper::AccountMappings::const_iterator iter = mappings.begin();
   EXPECT_EQ(mapping.email, iter->email);
   EXPECT_EQ(mapping.account_id, iter->account_id);
   EXPECT_TRUE(iter->access_token.empty());
   EXPECT_EQ(AccountMapping::REMOVING, iter->status);
   EXPECT_EQ(status_change_timestamp, iter->status_change_timestamp);
   EXPECT_TRUE(iter->last_message_id.empty());
 }
 
 // Tests that, if a new token arrives when the adding message is in progress
 // no new message is sent and account mapper still waits for the first one to
 // complete.
 TEST_F(GCMAccountMapperTest, TokenIsRefreshedWhenAdding) {
-  mapper()->Initialize(GCMAccountMapper::AccountMappings(), kRegistrationId);
+  mapper()->Initialize(GCMAccountMapper::AccountMappings());
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
 
   clock()->SetNow(base::Time::Now());
   std::vector<GCMClient::AccountTokenInfo> account_tokens;
   GCMClient::AccountTokenInfo account_token = MakeAccountTokenInfo("acc_id");
   account_tokens.push_back(account_token);
   mapper()->SetAccountTokens(account_tokens);
   DCHECK_EQ(CustomFakeGCMDriver::SEND_STARTED, gcm_driver().last_action());
 
   clock()->SetNow(base::Time::Now());
   gcm_driver().CompleteSend(gcm_driver().last_message_id(), GCMClient::SUCCESS);
   DCHECK_EQ(CustomFakeGCMDriver::SEND_FINISHED, gcm_driver().last_action());
 
   // Providing another token and clearing status.
   gcm_driver().Clear();
   mapper()->SetAccountTokens(account_tokens);
   DCHECK_EQ(CustomFakeGCMDriver::NONE, gcm_driver().last_action());
 }
 
 // Tests that, if a new token arrives when a removing message is in progress
 // a new adding message is sent and while account mapping status is changed to
 // mapped. If the original Removing message arrives it is discarded.
 TEST_F(GCMAccountMapperTest, TokenIsRefreshedWhenRemoving) {
   // Start with one account that is mapped.
   AccountMapping mapping = MakeAccountMapping(
       "acc_id", AccountMapping::MAPPED, base::Time::Now(), std::string());
 
   GCMAccountMapper::AccountMappings stored_mappings;
   stored_mappings.push_back(mapping);
-  mapper()->Initialize(stored_mappings, kRegistrationId);
+  mapper()->Initialize(stored_mappings);
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
   clock()->SetNow(base::Time::Now());
 
   // Remove the token to trigger a remove message to be sent
   mapper()->SetAccountTokens(std::vector<GCMClient::AccountTokenInfo>());
   EXPECT_EQ(CustomFakeGCMDriver::SEND_STARTED, gcm_driver().last_action());
   gcm_driver().CompleteSend(gcm_driver().last_message_id(), GCMClient::SUCCESS);
   EXPECT_EQ(CustomFakeGCMDriver::SEND_FINISHED, gcm_driver().last_action());
 
   std::string remove_message_id = gcm_driver().last_message_id();
   gcm_driver().Clear();
@@ skipping 32 matching lines @@
   clock()->SetNow(base::Time::Now());
   base::Time half_hour_ago = clock()->Now() - base::TimeDelta::FromMinutes(30);
   GCMAccountMapper::AccountMappings stored_mappings;
   stored_mappings.push_back(MakeAccountMapping(
       "acc_id_0", AccountMapping::ADDING, half_hour_ago, "acc_id_0_msg"));
   stored_mappings.push_back(MakeAccountMapping(
       "acc_id_1", AccountMapping::MAPPED, half_hour_ago, "acc_id_1_msg"));
   stored_mappings.push_back(MakeAccountMapping(
       "acc_id_2", AccountMapping::REMOVING, half_hour_ago, "acc_id_2_msg"));
 
-  mapper()->Initialize(stored_mappings, kRegistrationId);
+  mapper()->Initialize(stored_mappings);
+  gcm_driver().CompleteRegister(kRegistrationId, GCMClient::SUCCESS);
 
   GCMAccountMapper::AccountMappings expected_mappings(stored_mappings);
 
   // Finish messages after a restart.
   clock()->SetNow(base::Time::Now());
   gcm_driver().AcknowledgeSend(expected_mappings[0].last_message_id);
   expected_mappings[0].status_change_timestamp = clock()->Now();
   expected_mappings[0].status = AccountMapping::MAPPED;
   expected_mappings[0].last_message_id.clear();
 
@@ skipping 71 matching lines @@
   expected_mappings[1].status_change_timestamp = clock()->Now();
   expected_mappings[1].status = AccountMapping::MAPPED;
   expected_mappings[2].status_change_timestamp = clock()->Now();
   expected_mappings[2].status = AccountMapping::MAPPED;
 
   VerifyMappings(
       expected_mappings, GetAccounts(), "Step 6, Three mapped accounts");
 }
 
 }  // namespace gcm