[Sync] Use GCMDriver to deliver invalidations in tests

This CL uses GCMDriver directly to simulate sending sync invalidations
from the server. FakeGCMDriver also allows now to wait for specific GCMAppHandler to be added before connection (in real scenario GCMClient starts when the first AppHandler is added).

Note that Android has a different implementation and does not check for actual token to match. This should work since Android tests don't support multiple clients yet.

Bug: 1331206
Change-Id: I0dba63cdae469cb7047f076df1fd1a4d78153977
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4150488
Reviewed-by: Marc Treib <[email protected]>
Commit-Queue: Rushan Suleymanov <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1095004}
diff --git a/chrome/browser/login_detection/password_store_sites_browsertest.cc b/chrome/browser/login_detection/password_store_sites_browsertest.cc
index ab434d3..7f028a16 100644
--- a/chrome/browser/login_detection/password_store_sites_browsertest.cc
+++ b/chrome/browser/login_detection/password_store_sites_browsertest.cc
@@ -14,6 +14,7 @@
 #include "chrome/browser/ui/browser.h"
 #include "chrome/test/base/in_process_browser_test.h"
 #include "chrome/test/base/ui_test_utils.h"
+#include "components/keyed_service/content/browser_context_dependency_manager.h"
 #include "components/password_manager/core/browser/password_store.h"
 #include "components/password_manager/core/browser/password_store_interface.h"
 #include "components/password_manager/core/browser/test_password_store.h"
diff --git a/chrome/browser/sync/test/integration/device_info_helper.cc b/chrome/browser/sync/test/integration/device_info_helper.cc
index 7b80a74..8e7adc3 100644
--- a/chrome/browser/sync/test/integration/device_info_helper.cc
+++ b/chrome/browser/sync/test/integration/device_info_helper.cc
@@ -21,6 +21,7 @@
 }
 
 bool ServerDeviceInfoMatchChecker::IsExitConditionSatisfied(std::ostream* os) {
+  *os << "Waiting for server DeviceInfo to match: ";
   std::vector<sync_pb::SyncEntity> entities =
       fake_server()->GetSyncEntitiesByModelType(syncer::DEVICE_INFO);
 
diff --git a/chrome/browser/sync/test/integration/fake_sync_gcm_driver_for_instance_id.cc b/chrome/browser/sync/test/integration/fake_sync_gcm_driver_for_instance_id.cc
index bacddb3..431c7ce 100644
--- a/chrome/browser/sync/test/integration/fake_sync_gcm_driver_for_instance_id.cc
+++ b/chrome/browser/sync/test/integration/fake_sync_gcm_driver_for_instance_id.cc
@@ -4,39 +4,17 @@
 
 #include "chrome/browser/sync/test/integration/fake_sync_gcm_driver_for_instance_id.h"
 
+#include "base/files/file_path.h"
 #include "base/task/sequenced_task_runner.h"
-#include "base/task/single_thread_task_runner.h"
-#include "base/task/task_traits.h"
-#include "base/task/thread_pool.h"
 #include "chrome/browser/profiles/profile.h"
 #include "components/gcm_driver/crypto/gcm_encryption_result.h"
-#include "components/gcm_driver/fake_gcm_profile_service.h"
-#include "components/gcm_driver/gcm_profile_service.h"
-#include "components/keyed_service/core/keyed_service.h"
-
-// static
-std::unique_ptr<KeyedService> FakeSyncGCMDriver::Build(
-    content::BrowserContext* context) {
-  auto service = std::make_unique<gcm::FakeGCMProfileService>();
-  Profile* profile = Profile::FromBrowserContext(context);
-
-  // Allow blocking to initialize GCM client from the disk.
-  scoped_refptr<base::SequencedTaskRunner> blocking_task_runner(
-      base::ThreadPool::CreateSequencedTaskRunner(
-          {base::MayBlock(), base::TaskPriority::BEST_EFFORT,
-           base::TaskShutdownBehavior::SKIP_ON_SHUTDOWN}));
-  service->SetDriverForTesting(
-      std::make_unique<FakeSyncGCMDriver>(profile, blocking_task_runner));
-  return service;
-}
 
 FakeSyncGCMDriver::FakeSyncGCMDriver(
     Profile* profile,
     const scoped_refptr<base::SequencedTaskRunner>& blocking_task_runner)
     : instance_id::FakeGCMDriverForInstanceID(
           profile->GetPath().Append(FILE_PATH_LITERAL("gcm_test_store")),
-          blocking_task_runner),
-      profile_(profile) {}
+          blocking_task_runner) {}
 
 void FakeSyncGCMDriver::EncryptMessage(const std::string& app_id,
                                        const std::string& authorized_entity,
diff --git a/chrome/browser/sync/test/integration/fake_sync_gcm_driver_for_instance_id.h b/chrome/browser/sync/test/integration/fake_sync_gcm_driver_for_instance_id.h
index 3bc12684..906e16b6 100644
--- a/chrome/browser/sync/test/integration/fake_sync_gcm_driver_for_instance_id.h
+++ b/chrome/browser/sync/test/integration/fake_sync_gcm_driver_for_instance_id.h
@@ -8,29 +8,21 @@
 #include <memory>
 #include <string>
 
-#include "base/memory/raw_ptr.h"
 #include "base/memory/scoped_refptr.h"
 #include "components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.h"
 
+class Profile;
+
 namespace base {
 class SequencedTaskRunner;
 }  // namespace base
 
-namespace content {
-class BrowserContext;
-}  // namespace content
-
-class KeyedService;
-class Profile;
-
 class FakeSyncGCMDriver : public instance_id::FakeGCMDriverForInstanceID {
  public:
   FakeSyncGCMDriver(
       Profile* profile,
       const scoped_refptr<base::SequencedTaskRunner>& blocking_task_runner);
 
-  static std::unique_ptr<KeyedService> Build(content::BrowserContext* context);
-
  protected:
   // FakeGCMDriverForInstanceID overrides:
   void EncryptMessage(const std::string& app_id,
@@ -39,9 +31,6 @@
                       const std::string& auth_secret,
                       const std::string& message,
                       EncryptMessageCallback callback) override;
-
- private:
-  raw_ptr<Profile> profile_;
 };
 
 #endif  // CHROME_BROWSER_SYNC_TEST_INTEGRATION_FAKE_SYNC_GCM_DRIVER_FOR_INSTANCE_ID_H_
diff --git a/chrome/browser/sync/test/integration/invalidations/fake_server_sync_invalidation_sender.cc b/chrome/browser/sync/test/integration/invalidations/fake_server_sync_invalidation_sender.cc
index 5d25b78..67b75954d 100644
--- a/chrome/browser/sync/test/integration/invalidations/fake_server_sync_invalidation_sender.cc
+++ b/chrome/browser/sync/test/integration/invalidations/fake_server_sync_invalidation_sender.cc
@@ -4,21 +4,14 @@
 
 #include "chrome/browser/sync/test/integration/invalidations/fake_server_sync_invalidation_sender.h"
 
-#include "base/containers/contains.h"
+#include "base/containers/cxx20_erase.h"
+#include "base/logging.h"
 #include "base/time/time.h"
+#include "components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.h"
 #include "components/sync/base/time.h"
-#include "components/sync/invalidations/fcm_handler.h"
 
 namespace fake_server {
 
-namespace {
-
-// This has the same value as in
-// components/sync/invalidations/sync_invalidations_service_impl.cc.
-const char kSyncInvalidationsAppId[] = "com.google.chrome.sync.invalidations";
-
-}  // namespace
-
 FakeServerSyncInvalidationSender::FakeServerSyncInvalidationSender(
     FakeServer* fake_server)
     : fake_server_(fake_server) {
@@ -28,30 +21,23 @@
 
 FakeServerSyncInvalidationSender::~FakeServerSyncInvalidationSender() {
   fake_server_->RemoveObserver(this);
-
-  // Unsubscribe from all the remaining FCM handlers. This is mostly the case
-  // for Android platform.
-  for (syncer::FCMHandler* fcm_handler : fcm_handlers_) {
-    fcm_handler->RemoveTokenObserver(this);
+  for (instance_id::FakeGCMDriverForInstanceID* fake_gcm_driver :
+       fake_gcm_drivers_) {
+    fake_gcm_driver->RemoveConnectionObserver(this);
   }
 }
 
-void FakeServerSyncInvalidationSender::AddFCMHandler(
-    syncer::FCMHandler* fcm_handler) {
-  DCHECK(fcm_handler);
-  DCHECK(!base::Contains(fcm_handlers_, fcm_handler));
-
-  fcm_handlers_.push_back(fcm_handler);
-  fcm_handler->AddTokenObserver(this);
+void FakeServerSyncInvalidationSender::AddFakeGCMDriver(
+    instance_id::FakeGCMDriverForInstanceID* fake_gcm_driver) {
+  // It's safe to cast since SyncTest uses FakeGCMProfileService.
+  fake_gcm_drivers_.push_back(fake_gcm_driver);
+  fake_gcm_driver->AddConnectionObserver(this);
 }
 
-void FakeServerSyncInvalidationSender::RemoveFCMHandler(
-    syncer::FCMHandler* fcm_handler) {
-  DCHECK(fcm_handler);
-  DCHECK(base::Contains(fcm_handlers_, fcm_handler));
-
-  fcm_handler->RemoveTokenObserver(this);
-  base::Erase(fcm_handlers_, fcm_handler);
+void FakeServerSyncInvalidationSender::RemoveFakeGCMDriver(
+    instance_id::FakeGCMDriverForInstanceID* fake_gcm_driver) {
+  fake_gcm_driver->RemoveConnectionObserver(this);
+  base::Erase(fake_gcm_drivers_, fake_gcm_driver);
 }
 
 void FakeServerSyncInvalidationSender::OnWillCommit() {
@@ -92,18 +78,28 @@
   DeliverInvalidationsToHandlers();
 }
 
-void FakeServerSyncInvalidationSender::OnFCMRegistrationTokenChanged() {
+void FakeServerSyncInvalidationSender::OnConnected(
+    const net::IPEndPoint& ip_endpoint) {
+  // Try to deliver invalidations once GCMDriver is connected.
+  DVLOG(1) << "GCM driver connected";
   DeliverInvalidationsToHandlers();
 }
 
 void FakeServerSyncInvalidationSender::DeliverInvalidationsToHandlers() {
-  for (auto& token_and_invalidations : invalidations_to_deliver_) {
+  DVLOG(1) << "Trying to deliver invalidations for "
+           << invalidations_to_deliver_.size()
+           << " FCM tokens. Known target tokens from DeviceInfo: "
+           << token_to_interested_data_types_.size();
+  std::set<std::string> processed_tokens;
+  for (const auto& token_and_invalidations : invalidations_to_deliver_) {
     const std::string& token = token_and_invalidations.first;
-    // Pass a message to each FCMHandler to simulate a message from the
-    // GCMDriver.
+
+    // Pass a message to GCMDriver to simulate a message from the server.
     // TODO(crbug.com/1082115): Implement reflection blocking.
-    syncer::FCMHandler* fcm_handler = GetFCMHandlerByToken(token);
-    if (!fcm_handler) {
+    instance_id::FakeGCMDriverForInstanceID* fake_gcm_driver =
+        GetFakeGCMDriverByToken(token);
+    if (!fake_gcm_driver) {
+      DVLOG(1) << "Could not find FakeGCMDriver for token: " << token;
       continue;
     }
 
@@ -111,18 +107,37 @@
          token_and_invalidations.second) {
       gcm::IncomingMessage message;
       message.raw_data = payload.SerializeAsString();
-      fcm_handler->OnMessage(kSyncInvalidationsAppId, message);
+      fake_gcm_driver->DispatchMessage(kSyncInvalidationsAppId, message);
     }
 
-    token_and_invalidations.second.clear();
+    processed_tokens.insert(token);
+  }
+
+  for (const std::string& token_to_remove : processed_tokens) {
+    invalidations_to_deliver_.erase(token_to_remove);
   }
 }
 
-syncer::FCMHandler* FakeServerSyncInvalidationSender::GetFCMHandlerByToken(
+instance_id::FakeGCMDriverForInstanceID*
+FakeServerSyncInvalidationSender::GetFakeGCMDriverByToken(
     const std::string& fcm_registration_token) const {
-  for (syncer::FCMHandler* fcm_handler : fcm_handlers_) {
-    if (fcm_registration_token == fcm_handler->GetFCMRegistrationToken()) {
-      return fcm_handler;
+  for (instance_id::FakeGCMDriverForInstanceID* fake_gcm_driver :
+       fake_gcm_drivers_) {
+#if !BUILDFLAG(IS_ANDROID)
+    // On Android platform FCM registration token is returned from Java
+    // implementation, so HasTokenForAppId() does not contain these tokens.
+    // Since Android does not support several profiles, for the simplicity just
+    // check for AppHandler registration.
+    if (!fake_gcm_driver->HasTokenForAppId(kSyncInvalidationsAppId,
+                                           fcm_registration_token)) {
+      continue;
+    }
+#endif  // !BUILDFLAG(IS_ANDROID)
+
+    // AppHandler may not be registered while SyncSetup() is not called yet, the
+    // server should keep invalidations to deliver them later.
+    if (fake_gcm_driver->GetAppHandler(kSyncInvalidationsAppId)) {
+      return fake_gcm_driver;
     }
   }
   return nullptr;
diff --git a/chrome/browser/sync/test/integration/invalidations/fake_server_sync_invalidation_sender.h b/chrome/browser/sync/test/integration/invalidations/fake_server_sync_invalidation_sender.h
index f06f61c..c9acf91 100644
--- a/chrome/browser/sync/test/integration/invalidations/fake_server_sync_invalidation_sender.h
+++ b/chrome/browser/sync/test/integration/invalidations/fake_server_sync_invalidation_sender.h
@@ -6,24 +6,28 @@
 #define CHROME_BROWSER_SYNC_TEST_INTEGRATION_INVALIDATIONS_FAKE_SERVER_SYNC_INVALIDATION_SENDER_H_
 
 #include "base/memory/raw_ptr.h"
+#include "components/gcm_driver/gcm_connection_observer.h"
 #include "components/sync/base/model_type.h"
-#include "components/sync/invalidations/fcm_registration_token_observer.h"
 #include "components/sync/protocol/sync_invalidations_payload.pb.h"
 #include "components/sync/test/fake_server.h"
 
-namespace syncer {
-class FCMHandler;
-}
+namespace instance_id {
+class FakeGCMDriverForInstanceID;
+}  // namespace instance_id
 
 namespace fake_server {
 
 // This class is observing changes to the fake server, and sends invalidations
 // to clients upon commits. Sent invalidation follows the same format expected
 // by the sync invalidations framework (i.e. SyncInvalidationsService).
-class FakeServerSyncInvalidationSender
-    : public FakeServer::Observer,
-      public syncer::FCMRegistrationTokenObserver {
+class FakeServerSyncInvalidationSender : public FakeServer::Observer,
+                                         public gcm::GCMConnectionObserver {
  public:
+  // This has the same value as in
+  // components/sync/invalidations/sync_invalidations_service_impl.cc.
+  static constexpr char kSyncInvalidationsAppId[] =
+      "com.google.chrome.sync.invalidations";
+
   // |fake_server| must not be nullptr, and must outlive this object.
   explicit FakeServerSyncInvalidationSender(FakeServer* fake_server);
   ~FakeServerSyncInvalidationSender() override;
@@ -32,26 +36,24 @@
   FakeServerSyncInvalidationSender& operator=(
       const FakeServerSyncInvalidationSender&) = delete;
 
-  // |fcm_handler| must not be nullptr, and must be removed using
-  // RemoveFCMHandler(). If the FCM handler has registered token, all the
-  // message for the token will be delivered immediately.
-  void AddFCMHandler(syncer::FCMHandler* fcm_handler);
+  // Add |fake_gcm_driver| to send invalidations to from the fake server.
+  void AddFakeGCMDriver(
+      instance_id::FakeGCMDriverForInstanceID* fake_gcm_driver);
 
-  // |fcm_handler| must not be nullptr, and must exist in |fcm_handlers_|.
-  void RemoveFCMHandler(syncer::FCMHandler* fcm_handler);
+  // Remove |fake_gcm_driver| to stop sending invalidations.
+  void RemoveFakeGCMDriver(
+      instance_id::FakeGCMDriverForInstanceID* fake_gcm_driver);
 
   // FakeServer::Observer implementation.
   void OnWillCommit() override;
   void OnCommit(const std::string& committer_invalidator_client_id,
                 syncer::ModelTypeSet committed_model_types) override;
 
-  // syncer::FCMRegistrationTokenObserver implementation.
-  void OnFCMRegistrationTokenChanged() override;
+  // gcm::GCMConnectionObserver implementation.
+  void OnConnected(const net::IPEndPoint& ip_endpoint) override;
 
  private:
-  // Returns a corresponding FCM handler having the same
-  // |fcm_registration_token| if exists. Otherwise, returns nullptr.
-  syncer::FCMHandler* GetFCMHandlerByToken(
+  instance_id::FakeGCMDriverForInstanceID* GetFakeGCMDriverByToken(
       const std::string& fcm_registration_token) const;
 
   // Delivers all the incoming messages to the corresponding FCM handlers.
@@ -63,7 +65,6 @@
   void UpdateTokenToInterestedDataTypesMap();
 
   raw_ptr<FakeServer> fake_server_;
-  std::vector<syncer::FCMHandler*> fcm_handlers_;
 
   // Cache of invalidations to be dispatched by
   // DeliverInvalidationsToHandlers(), keyed by FCM registration token. If no
@@ -73,8 +74,11 @@
       invalidations_to_deliver_;
 
   // List of tokens with a list of interested data types. Used to send
-  // invalidations to a corresponding FCMHandler.
+  // invalidations to a corresponding client.
   std::map<std::string, syncer::ModelTypeSet> token_to_interested_data_types_;
+
+  std::vector<base::raw_ptr<instance_id::FakeGCMDriverForInstanceID>>
+      fake_gcm_drivers_;
 };
 
 }  // namespace fake_server
diff --git a/chrome/browser/sync/test/integration/single_client_bookmarks_sync_test.cc b/chrome/browser/sync/test/integration/single_client_bookmarks_sync_test.cc
index d5aa602f..c141c683 100644
--- a/chrome/browser/sync/test/integration/single_client_bookmarks_sync_test.cc
+++ b/chrome/browser/sync/test/integration/single_client_bookmarks_sync_test.cc
@@ -29,6 +29,7 @@
 #include "components/bookmarks/browser/bookmark_model.h"
 #include "components/bookmarks/browser/bookmark_node.h"
 #include "components/bookmarks/browser/url_and_title.h"
+#include "components/keyed_service/content/browser_context_dependency_manager.h"
 #include "components/keyed_service/core/keyed_service.h"
 #include "components/sync/base/client_tag_hash.h"
 #include "components/sync/base/command_line_switches.h"
diff --git a/chrome/browser/sync/test/integration/sync_test.cc b/chrome/browser/sync/test/integration/sync_test.cc
index 5491fa1..2c446633 100644
--- a/chrome/browser/sync/test/integration/sync_test.cc
+++ b/chrome/browser/sync/test/integration/sync_test.cc
@@ -21,6 +21,9 @@
 #include "base/strings/string_util.h"
 #include "base/strings/stringprintf.h"
 #include "base/strings/utf_string_conversions.h"
+#include "base/task/sequenced_task_runner.h"
+#include "base/task/task_traits.h"
+#include "base/task/thread_pool.h"
 #include "base/test/test_timeouts.h"
 #include "build/build_config.h"
 #include "build/chromeos_buildflags.h"
@@ -36,7 +39,6 @@
 #include "chrome/browser/profiles/profiles_state.h"
 #include "chrome/browser/signin/chrome_signin_client_factory.h"
 #include "chrome/browser/signin/identity_manager_factory.h"
-#include "chrome/browser/sync/sync_invalidations_service_factory.h"
 #include "chrome/browser/sync/sync_service_factory.h"
 #include "chrome/browser/sync/test/integration/committed_all_nudged_changes_checker.h"
 #include "chrome/browser/sync/test/integration/device_info_helper.h"
@@ -64,6 +66,7 @@
 #include "components/invalidation/impl/profile_identity_provider.h"
 #include "components/invalidation/impl/profile_invalidation_provider.h"
 #include "components/invalidation/public/invalidation_service.h"
+#include "components/keyed_service/content/browser_context_dependency_manager.h"
 #include "components/keyed_service/core/keyed_service.h"
 #include "components/os_crypt/os_crypt_mocker.h"
 #include "components/prefs/scoped_user_pref_update.h"
@@ -75,7 +78,6 @@
 #include "components/sync/driver/sync_service_impl.h"
 #include "components/sync/driver/sync_user_settings.h"
 #include "components/sync/engine/sync_scheduler_impl.h"
-#include "components/sync/invalidations/fcm_handler.h"
 #include "components/sync/invalidations/sync_invalidations_service_impl.h"
 #include "components/sync/test/fake_server_network_resources.h"
 #include "content/public/browser/navigation_entry.h"
@@ -208,6 +210,29 @@
   return it != profile_to_fcm_network_handler_map->end() ? it->second : nullptr;
 }
 
+std::unique_ptr<KeyedService> CreateGCMProfileService(
+    content::BrowserContext* context) {
+  scoped_refptr<base::SequencedTaskRunner> blocking_task_runner(
+      base::ThreadPool::CreateSequencedTaskRunner(
+          {base::MayBlock(), base::TaskPriority::BEST_EFFORT,
+           base::TaskShutdownBehavior::SKIP_ON_SHUTDOWN}));
+
+  Profile* profile = Profile::FromBrowserContext(context);
+  auto service = std::make_unique<gcm::FakeGCMProfileService>(
+      std::make_unique<FakeSyncGCMDriver>(profile, blocking_task_runner));
+
+  service->GetFakeGCMDriver()->WaitForAppIdBeforeConnection(
+      fake_server::FakeServerSyncInvalidationSender::kSyncInvalidationsAppId);
+  return service;
+}
+
+instance_id::FakeGCMDriverForInstanceID* GetFakeGCMDriverForProfile(
+    Profile* profile) {
+  return static_cast<gcm::FakeGCMProfileService*>(
+             gcm::GCMProfileServiceFactory::GetForProfile(profile))
+      ->GetFakeGCMDriver();
+}
+
 }  // namespace
 
 #if !BUILDFLAG(IS_ANDROID)
@@ -665,17 +690,8 @@
     sync_service_impl->OverrideNetworkForTest(
         fake_server::CreateFakeServerHttpPostProviderFactory(
             GetFakeServer()->AsWeakPtr()));
-    // TODO(crbug.com/1331206): use GCM driver directly to deliver
-    // invalidations.
-    syncer::SyncInvalidationsServiceImpl* sync_invalidations_service =
-        static_cast<syncer::SyncInvalidationsServiceImpl*>(
-            SyncInvalidationsServiceFactory::GetForProfile(profile));
-    if (sync_invalidations_service) {
-      profile_to_fcm_handler_map_[profile] =
-          sync_invalidations_service->GetFCMHandlerForTesting();
-      fake_server_sync_invalidation_sender_->AddFCMHandler(
-          sync_invalidations_service->GetFCMHandlerForTesting());
-    }
+    fake_server_sync_invalidation_sender_->AddFakeGCMDriver(
+        GetFakeGCMDriverForProfile(profile));
   }
 
   SyncServiceImplHarness::SigninType signin_type =
@@ -907,7 +923,6 @@
              observer : fake_server_invalidation_observers_) {
       fake_server_->RemoveObserver(observer.get());
     }
-    profile_to_fcm_handler_map_.clear();
     fake_server_sync_invalidation_sender_.reset();
     fake_server_.reset();
   }
@@ -962,19 +977,14 @@
     }
 
     CheckForDataTypeFailures(/*client_index=*/index);
+    fake_server_sync_invalidation_sender_->RemoveFakeGCMDriver(
+        GetFakeGCMDriverForProfile(profile));
     profiles_[index] = nullptr;
     clients_[index].reset();
 #if !BUILDFLAG(IS_ANDROID)
     DCHECK(!browsers_[index]);
 #endif  // !BUILDFLAG(IS_ANDROID)
   }
-
-  if (fake_server_sync_invalidation_sender_) {
-    DCHECK(base::Contains(profile_to_fcm_handler_map_, profile));
-    fake_server_sync_invalidation_sender_->RemoveFCMHandler(
-        profile_to_fcm_handler_map_[profile]);
-    profile_to_fcm_handler_map_.erase(profile);
-  }
 }
 
 void SyncTest::OnWillCreateBrowserContextServices(
@@ -991,7 +1001,7 @@
           base::BindRepeating(&SyncTest::CreateProfileInvalidationProvider,
                               &profile_to_fcm_network_handler_map_));
   gcm::GCMProfileServiceFactory::GetInstance()->SetTestingFactory(
-      context, base::BindRepeating(&FakeSyncGCMDriver::Build));
+      context, base::BindRepeating(&CreateGCMProfileService));
 }
 
 // static
@@ -1104,12 +1114,6 @@
           std::make_unique<fake_server::FakeServerSyncInvalidationSender>(
               fake_server_.get());
 
-      // Subscribe to invalidations for all the profiles which were created
-      // before. This is mainly the case on Android platform.
-      for (const auto& profile_and_fcm_handler : profile_to_fcm_handler_map_) {
-        fake_server_sync_invalidation_sender_->AddFCMHandler(
-            profile_and_fcm_handler.second);
-      }
       SetupMockGaiaResponses();
       SetupMockGaiaResponsesForProfile(
           ProfileManager::GetLastUsedProfileIfLoaded());
diff --git a/chrome/browser/sync/test/integration/sync_test.h b/chrome/browser/sync/test/integration/sync_test.h
index e3f3471..a0e4f43 100644
--- a/chrome/browser/sync/test/integration/sync_test.h
+++ b/chrome/browser/sync/test/integration/sync_test.h
@@ -23,7 +23,6 @@
 #include "chrome/browser/sync/test/integration/fake_server_invalidation_sender.h"
 #include "chrome/browser/sync/test/integration/invalidations/fake_server_sync_invalidation_sender.h"
 #include "chrome/common/buildflags.h"
-#include "components/keyed_service/content/browser_context_dependency_manager.h"
 #include "components/sync/base/model_type.h"
 #include "components/sync/base/user_selectable_type.h"
 #include "components/sync/test/fake_server.h"
@@ -75,7 +74,6 @@
 }  // namespace fake_server
 
 namespace syncer {
-class FCMHandler;
 class SyncServiceImpl;
 }  // namespace syncer
 
@@ -446,8 +444,6 @@
   std::map<const Profile*, invalidation::FCMNetworkHandler*>
       profile_to_fcm_network_handler_map_;
 
-  std::map<const Profile*, syncer::FCMHandler*> profile_to_fcm_handler_map_;
-
   // Triggers a GetUpdates via refresh after a configuration.
   std::unique_ptr<ConfigurationRefresher> configuration_refresher_;
 
diff --git a/chrome/browser/sync/test/integration/updated_progress_marker_checker.cc b/chrome/browser/sync/test/integration/updated_progress_marker_checker.cc
index 0cc6417..fb237fa 100644
--- a/chrome/browser/sync/test/integration/updated_progress_marker_checker.cc
+++ b/chrome/browser/sync/test/integration/updated_progress_marker_checker.cc
@@ -26,9 +26,10 @@
 UpdatedProgressMarkerChecker::~UpdatedProgressMarkerChecker() = default;
 
 bool UpdatedProgressMarkerChecker::IsExitConditionSatisfied(std::ostream* os) {
-  *os << "Waiting for progress markers";
+  *os << "Waiting for progress markers... ";
 
   if (!has_unsynced_items_.has_value()) {
+    *os << "Unknown synced values state.";
     return false;
   }
 
@@ -42,9 +43,22 @@
   //    by the test-only 'self-notify' cycle).
   // 3. No pending local changes (which will ultimately generate new progress
   //    markers once submitted to the server).
-  return !snap.download_progress_markers().empty() &&
-         snap.model_neutral_state().num_successful_commits == 0 &&
-         !has_unsynced_items_.value();
+  if (snap.download_progress_markers().empty()) {
+    *os << "Progress markers are empty.";
+    return false;
+  }
+
+  if (snap.model_neutral_state().num_successful_commits > 0) {
+    *os << "Last sync cycle wasn't empty.";
+    return false;
+  }
+
+  if (has_unsynced_items_.value()) {
+    *os << "Has unsynced items.";
+    return false;
+  }
+
+  return true;
 }
 
 void UpdatedProgressMarkerChecker::GotHasUnsyncedItems(
diff --git a/components/gcm_driver/fake_gcm_profile_service.cc b/components/gcm_driver/fake_gcm_profile_service.cc
index 6437627..fa81e44 100644
--- a/components/gcm_driver/fake_gcm_profile_service.cc
+++ b/components/gcm_driver/fake_gcm_profile_service.cc
@@ -11,12 +11,10 @@
 #include "base/location.h"
 #include "base/memory/raw_ptr.h"
 #include "base/memory/weak_ptr.h"
-#include "base/strings/string_number_conversions.h"
 #include "base/strings/stringprintf.h"
 #include "base/task/single_thread_task_runner.h"
 #include "base/time/time.h"
 #include "components/gcm_driver/crypto/gcm_encryption_result.h"
-#include "components/gcm_driver/fake_gcm_client_factory.h"
 #include "components/gcm_driver/gcm_driver.h"
 #include "components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.h"
 
@@ -25,7 +23,10 @@
 class FakeGCMProfileService::CustomFakeGCMDriver
     : public instance_id::FakeGCMDriverForInstanceID {
  public:
-  explicit CustomFakeGCMDriver(FakeGCMProfileService* service);
+  CustomFakeGCMDriver();
+
+  // Must be called before any other methods.
+  void SetService(FakeGCMProfileService* service);
 
   CustomFakeGCMDriver(const CustomFakeGCMDriver&) = delete;
   CustomFakeGCMDriver& operator=(const CustomFakeGCMDriver&) = delete;
@@ -39,9 +40,6 @@
                       const std::string& message_id,
                       GCMClient::Result result);
 
-  void OnDispatchMessage(const std::string& app_id,
-                         const IncomingMessage& message);
-
   // GCMDriver overrides:
   void EncryptMessage(const std::string& app_id,
                       const std::string& authorized_entity,
@@ -80,7 +78,7 @@
               const std::string& receiver_id,
               const OutgoingMessage& message);
 
-  raw_ptr<FakeGCMProfileService> service_;
+  raw_ptr<FakeGCMProfileService> service_ = nullptr;
 
   // Used to give each registration a unique registration id. Does not decrease
   // when unregister is called.
@@ -90,11 +88,9 @@
       this};  // Must be last.
 };
 
-FakeGCMProfileService::CustomFakeGCMDriver::CustomFakeGCMDriver(
-    FakeGCMProfileService* service)
-    : service_(service) {}
+FakeGCMProfileService::CustomFakeGCMDriver::CustomFakeGCMDriver() = default;
 
-FakeGCMProfileService::CustomFakeGCMDriver::~CustomFakeGCMDriver() {}
+FakeGCMProfileService::CustomFakeGCMDriver::~CustomFakeGCMDriver() = default;
 
 void FakeGCMProfileService::CustomFakeGCMDriver::RegisterImpl(
     const std::string& app_id,
@@ -171,6 +167,11 @@
   std::move(callback).Run(GCMEncryptionResult::ENCRYPTED_DRAFT_08, message);
 }
 
+void FakeGCMProfileService::CustomFakeGCMDriver::SetService(
+    FakeGCMProfileService* service) {
+  service_ = service;
+}
+
 void FakeGCMProfileService::CustomFakeGCMDriver::DoSend(
     const std::string& app_id,
     const std::string& receiver_id,
@@ -207,24 +208,22 @@
       app_id, authorized_entity, scope, std::move(callback));
 }
 
-void FakeGCMProfileService::CustomFakeGCMDriver::OnDispatchMessage(
-    const std::string& app_id,
-    const IncomingMessage& message) {
-  DispatchMessage(app_id, message);
-}
-
 // static
 std::unique_ptr<KeyedService> FakeGCMProfileService::Build(
     content::BrowserContext* context) {
-  std::unique_ptr<FakeGCMProfileService> service =
-      std::make_unique<FakeGCMProfileService>();
-  service->SetDriverForTesting(
-      std::make_unique<CustomFakeGCMDriver>(service.get()));
+  auto custom_driver = std::make_unique<CustomFakeGCMDriver>();
+  CustomFakeGCMDriver* custom_driver_ptr = custom_driver.get();
 
+  std::unique_ptr<FakeGCMProfileService> service =
+      std::make_unique<FakeGCMProfileService>(std::move(custom_driver));
+
+  custom_driver_ptr->SetService(service.get());
   return service;
 }
 
-FakeGCMProfileService::FakeGCMProfileService() = default;
+FakeGCMProfileService::FakeGCMProfileService(
+    std::unique_ptr<instance_id::FakeGCMDriverForInstanceID> fake_gcm_driver)
+    : GCMProfileService(std::move(fake_gcm_driver)) {}
 
 FakeGCMProfileService::~FakeGCMProfileService() = default;
 
@@ -235,9 +234,12 @@
 
 void FakeGCMProfileService::DispatchMessage(const std::string& app_id,
                                             const IncomingMessage& message) {
-  CustomFakeGCMDriver* custom_driver =
-      static_cast<CustomFakeGCMDriver*>(driver());
-  custom_driver->OnDispatchMessage(app_id, message);
+  GetFakeGCMDriver()->DispatchMessage(app_id, message);
+}
+
+instance_id::FakeGCMDriverForInstanceID*
+FakeGCMProfileService::GetFakeGCMDriver() {
+  return static_cast<instance_id::FakeGCMDriverForInstanceID*>(driver());
 }
 
 }  // namespace gcm
diff --git a/components/gcm_driver/fake_gcm_profile_service.h b/components/gcm_driver/fake_gcm_profile_service.h
index d53777c..0adf038 100644
--- a/components/gcm_driver/fake_gcm_profile_service.h
+++ b/components/gcm_driver/fake_gcm_profile_service.h
@@ -18,6 +18,10 @@
 class BrowserContext;
 }  // namespace content
 
+namespace instance_id {
+class FakeGCMDriverForInstanceID;
+}  // namespace instance_id
+
 namespace gcm {
 
 // Acts as a bridge between GCM API and GCM Client layer for testing purposes.
@@ -26,7 +30,8 @@
   // Helper function to be used with KeyedServiceFactory::SetTestingFactory().
   static std::unique_ptr<KeyedService> Build(content::BrowserContext* context);
 
-  FakeGCMProfileService();
+  explicit FakeGCMProfileService(
+      std::unique_ptr<instance_id::FakeGCMDriverForInstanceID> fake_gcm_driver);
 
   FakeGCMProfileService(const FakeGCMProfileService&) = delete;
   FakeGCMProfileService& operator=(const FakeGCMProfileService&) = delete;
@@ -38,6 +43,8 @@
   void DispatchMessage(const std::string& app_id,
                        const IncomingMessage& message);
 
+  instance_id::FakeGCMDriverForInstanceID* GetFakeGCMDriver();
+
   const OutgoingMessage& last_sent_message() const {
     return last_sent_message_;
   }
diff --git a/components/gcm_driver/gcm_profile_service.cc b/components/gcm_driver/gcm_profile_service.cc
index addd858..0f2be5a 100644
--- a/components/gcm_driver/gcm_profile_service.cc
+++ b/components/gcm_driver/gcm_profile_service.cc
@@ -184,9 +184,17 @@
 }
 #endif  // BUILDFLAG(USE_GCM_FROM_PLATFORM)
 
-GCMProfileService::GCMProfileService() {}
+GCMProfileService::GCMProfileService(std::unique_ptr<GCMDriver> gcm_driver)
+    : driver_(std::move(gcm_driver)) {
+#if !BUILDFLAG(USE_GCM_FROM_PLATFORM)
+  if (identity_observer_) {
+    identity_observer_ = std::make_unique<IdentityObserver>(
+        identity_manager_, url_loader_factory_, driver_.get());
+  }
+#endif  // !BUILDFLAG(USE_GCM_FROM_PLATFORM)
+}
 
-GCMProfileService::~GCMProfileService() {}
+GCMProfileService::~GCMProfileService() = default;
 
 void GCMProfileService::Shutdown() {
 #if !BUILDFLAG(USE_GCM_FROM_PLATFORM)
@@ -198,15 +206,4 @@
   }
 }
 
-void GCMProfileService::SetDriverForTesting(std::unique_ptr<GCMDriver> driver) {
-  driver_ = std::move(driver);
-
-#if !BUILDFLAG(USE_GCM_FROM_PLATFORM)
-  if (identity_observer_) {
-    identity_observer_ = std::make_unique<IdentityObserver>(
-        identity_manager_, url_loader_factory_, driver.get());
-  }
-#endif  // !BUILDFLAG(USE_GCM_FROM_PLATFORM)
-}
-
 }  // namespace gcm
diff --git a/components/gcm_driver/gcm_profile_service.h b/components/gcm_driver/gcm_profile_service.h
index 81a282cc..aeead5c 100644
--- a/components/gcm_driver/gcm_profile_service.h
+++ b/components/gcm_driver/gcm_profile_service.h
@@ -78,14 +78,11 @@
   // KeyedService:
   void Shutdown() override;
 
-  // For testing purposes.
-  void SetDriverForTesting(std::unique_ptr<GCMDriver> driver);
-
   GCMDriver* driver() const { return driver_.get(); }
 
  protected:
   // Used for constructing fake GCMProfileService for testing purpose.
-  GCMProfileService();
+  explicit GCMProfileService(std::unique_ptr<GCMDriver> driver);
 
  private:
   std::unique_ptr<GCMDriver> driver_;
diff --git a/components/gcm_driver/instance_id/BUILD.gn b/components/gcm_driver/instance_id/BUILD.gn
index fc71330..f9bfee62f 100644
--- a/components/gcm_driver/instance_id/BUILD.gn
+++ b/components/gcm_driver/instance_id/BUILD.gn
@@ -53,6 +53,7 @@
   deps = [
     "//base",
     "//components/gcm_driver",
+    "//net",
     "//testing/gtest",
   ]
 
diff --git a/components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.cc b/components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.cc
index 22b64453..1a4be74 100644
--- a/components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.cc
+++ b/components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.cc
@@ -11,12 +11,18 @@
 #include "base/location.h"
 #include "base/logging.h"
 #include "base/rand_util.h"
+#include "base/ranges/algorithm.h"
 #include "base/strings/string_number_conversions.h"
 #include "base/strings/string_util.h"
 #include "base/task/sequenced_task_runner.h"
 #include "base/task/single_thread_task_runner.h"
+#include "base/task/task_runner.h"
 #include "base/threading/thread_restrictions.h"
+#include "base/time/time.h"
+#include "components/gcm_driver/common/gcm_message.h"
 #include "components/gcm_driver/gcm_client.h"
+#include "components/gcm_driver/gcm_connection_observer.h"
+#include "net/base/ip_endpoint.h"
 
 namespace instance_id {
 
@@ -53,10 +59,10 @@
       << "Failed to read data from stored FCM tokens file";
 
   for (const auto [key, value] : data.value().GetDict()) {
+    DVLOG(1) << "Loaded FCM token from file, key: " << key
+             << ", value: " << value.GetString();
     tokens_[key] = value.GetString();
   }
-
-  DVLOG(1) << "Loaded tokens from file: " << tokens_.size();
 }
 
 FakeGCMDriverForInstanceID::~FakeGCMDriverForInstanceID() = default;
@@ -66,6 +72,48 @@
   return this;
 }
 
+void FakeGCMDriverForInstanceID::AddConnectionObserver(
+    gcm::GCMConnectionObserver* observer) {
+  connection_observers_.AddObserver(observer);
+}
+
+void FakeGCMDriverForInstanceID::RemoveConnectionObserver(
+    gcm::GCMConnectionObserver* observer) {
+  connection_observers_.RemoveObserver(observer);
+}
+
+void FakeGCMDriverForInstanceID::AddAppHandler(const std::string& app_id,
+                                               gcm::GCMAppHandler* handler) {
+  FakeGCMDriver::AddAppHandler(app_id, handler);
+
+  DVLOG(1) << "GCMAppHandler was added: " << app_id;
+
+  if (app_id_for_connection_.empty() || app_id == app_id_for_connection_) {
+    ConnectIfNeeded();
+  }
+}
+
+bool FakeGCMDriverForInstanceID::HasTokenForAppId(
+    const std::string& app_id,
+    const std::string& token) const {
+#if BUILDFLAG(IS_ANDROID)
+  // FCM registration tokens on Android should be handled by
+  // FakeInstanceIDWithSubtype.
+  NOTREACHED();
+#endif  // BUILDFLAG(IS_ANDROID)
+  for (const auto& [key, stored_token] : tokens_) {
+    if (token == stored_token && base::StartsWith(key, app_id)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+void FakeGCMDriverForInstanceID::WaitForAppIdBeforeConnection(
+    const std::string& app_id) {
+  app_id_for_connection_ = app_id;
+}
+
 void FakeGCMDriverForInstanceID::AddInstanceIDData(
     const std::string& app_id,
     const std::string& instance_id,
@@ -183,4 +231,16 @@
   DCHECK(success) << "Failed to store FCM tokens";
 }
 
+void FakeGCMDriverForInstanceID::ConnectIfNeeded() {
+  if (connected_) {
+    return;
+  }
+
+  DVLOG(1) << "GCMDriver connected.";
+  connected_ = true;
+  for (gcm::GCMConnectionObserver& observer : connection_observers_) {
+    observer.OnConnected(net::IPEndPoint());
+  }
+}
+
 }  // namespace instance_id
diff --git a/components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.h b/components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.h
index 8e2f7bf..d3263a8a 100644
--- a/components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.h
+++ b/components/gcm_driver/instance_id/fake_gcm_driver_for_instance_id.h
@@ -10,6 +10,8 @@
 #include <utility>
 
 #include "base/files/file_path.h"
+#include "base/memory/weak_ptr.h"
+#include "base/observer_list.h"
 #include "components/gcm_driver/fake_gcm_driver.h"
 
 namespace base {
@@ -34,6 +36,21 @@
 
   // FakeGCMDriver overrides:
   gcm::InstanceIDHandler* GetInstanceIDHandlerInternal() override;
+  void AddConnectionObserver(gcm::GCMConnectionObserver* observer) override;
+  void RemoveConnectionObserver(gcm::GCMConnectionObserver* observer) override;
+  void AddAppHandler(const std::string& app_id,
+                     gcm::GCMAppHandler* handler) override;
+
+  // Expose protected method for testing.
+  using gcm::FakeGCMDriver::DispatchMessage;
+
+  // Returns true if the given |app_id| has the expected |token|. Note that
+  // tokens may be loaded before GCMDriver is connected.
+  bool HasTokenForAppId(const std::string& app_id,
+                        const std::string& token) const;
+
+  // GCMDriver will not connect until the given |app_id| is added.
+  void WaitForAppIdBeforeConnection(const std::string& app_id);
 
   const std::string& last_gettoken_app_id() const {
     return last_gettoken_app_id_;
@@ -69,8 +86,14 @@
                          GetInstanceIDDataCallback callback) override;
 
  private:
+  // Stores generated FCM registration tokens to a file, to keep the same tokens
+  // across browser restarts in tests.
   void StoreTokensIfNeeded();
 
+  // Used to simulate connection of GCMDriver after adding the first
+  // GCMAppHandler.
+  void ConnectIfNeeded();
+
   std::string GenerateTokenImpl(const std::string& app_id,
                                 const std::string& authorized_entity,
                                 const std::string& scope);
@@ -83,6 +106,20 @@
   std::string last_gettoken_app_id_;
   std::string last_gettoken_authorized_entity_;
   std::string last_deletetoken_app_id_;
+
+  // Simulate a connection to the server only after the given AppHandler has
+  // been added. This is required to prevent message loss in GCMDriver while
+  // dispatching a message.
+  // TODO(crbug.com/1408769): remove once GCMDriver fixes it.
+  std::string app_id_for_connection_;
+  bool connected_ = false;
+
+  base::ObserverList<gcm::GCMConnectionObserver,
+                     /*check_empty=*/false,
+                     /*allow_reentrancy=*/false>::Unchecked
+      connection_observers_;
+
+  base::WeakPtrFactory<FakeGCMDriverForInstanceID> weak_ptr_factory_{this};
 };
 
 }  // namespace instance_id