Fix flaky UKM browser tests

Bug: 996823
Change-Id: I74170c9ce3008e40c280d0b7ccb58ffc10793364
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1784878Reviewed-by: François Doray <fdoray@chromium.org>
Reviewed-by: Alexei Svitkine <asvitkine@chromium.org>
Commit-Queue: Yue Ru Sun <yrsun@chromium.org>
Cr-Commit-Position: refs/heads/master@{#695648}

chrome/browser/metrics/ukm_browsertest.cc

@@ -206,7 +206,12 @@ class UkmBrowserTestBase : public SyncTest {
   void BuildAndStoreUkmLog() {
     auto* service = ukm_service();
     DCHECK(service);
+    // Wait for initialization to complete before flushing.
+    base::RunLoop run_loop;
+    service->SetInitializationCompleteCallbackForTesting(run_loop.QuitClosure());
+    run_loop.Run();
     DCHECK(service->initialize_complete_);
+
     service->Flush();
     DCHECK(service->reporting_service_.ukm_log_store()->has_unsent_logs());
   }
@@ -540,13 +545,7 @@ IN_PROC_BROWSER_TEST_P(UkmBrowserTest, OpenNonSyncCheck) {
 // Keep in sync with UkmTest.testMetricConsent in
 // chrome/android/javatests/src/org/chromium/chrome/browser/sync/
 // UkmTest.java.
-#if defined(OS_CHROMEOS)
-// TODO(https://crbug.com/996823): Re-enable this test.
-#define MAYBE_MetricsConsentCheck DISABLED_MetricsConsentCheck
-#else
-#define MAYBE_MetricsConsentCheck MetricsConsentCheck
-#endif
-IN_PROC_BROWSER_TEST_P(UkmBrowserTest, MAYBE_MetricsConsentCheck) {
+IN_PROC_BROWSER_TEST_P(UkmBrowserTest, MetricsConsentCheck) {
   MetricsConsentOverride metrics_consent(true);
 
   Profile* profile = ProfileManager::GetActiveUserProfile();
@@ -576,13 +575,7 @@ IN_PROC_BROWSER_TEST_P(UkmBrowserTest, MAYBE_MetricsConsentCheck) {
   CloseBrowserSynchronously(sync_browser);
 }
 
-#if defined(OS_CHROMEOS)
-// TODO(https://crbug.com/996823): Re-enable this test.
-#define MAYBE_LogProtoData DISABLED_LogProtoData
-#else
-#define MAYBE_LogProtoData LogProtoData
-#endif
-IN_PROC_BROWSER_TEST_P(UkmBrowserTest, MAYBE_LogProtoData) {
+IN_PROC_BROWSER_TEST_P(UkmBrowserTest, LogProtoData) {
   MetricsConsentOverride metrics_consent(true);
 
   Profile* profile = ProfileManager::GetActiveUserProfile();
@@ -622,16 +615,7 @@ IN_PROC_BROWSER_TEST_P(UkmBrowserTest, MAYBE_LogProtoData) {
 
 // Verifies that network provider attaches effective connection type correctly
 // to the UKM report.
-#if defined(OS_CHROMEOS)
-// TODO(https://crbug.com/996823): Re-enable this test.
-#define MAYBE_NetworkProviderPopulatesSystemProfile \
-  DISABLED_NetworkProviderPopulatesSystemProfile
-#else
-#define MAYBE_NetworkProviderPopulatesSystemProfile \
-  NetworkProviderPopulatesSystemProfile
-#endif
-IN_PROC_BROWSER_TEST_P(UkmBrowserTest,
-                       MAYBE_NetworkProviderPopulatesSystemProfile) {
+IN_PROC_BROWSER_TEST_P(UkmBrowserTest, NetworkProviderPopulatesSystemProfile) {
   // Override network quality to 2G. This should cause the
   // |max_effective_connection_type| in the system profile to be set to 2G.
   g_browser_process->network_quality_tracker()

components/ukm/ukm_service.cc

@@ -225,6 +225,9 @@ void UkmService::FinishedInitTask() {
   DVLOG(1) << "UkmService::FinishedInitTask";
   initialize_complete_ = true;
   scheduler_->InitTaskComplete();
+  if (initialization_complete_callback_) {
+    std::move(initialization_complete_callback_).Run();
+  }
 }
 
 void UkmService::RotateLog() {
@@ -268,4 +271,13 @@ bool UkmService::ShouldRestrictToWhitelistedEntries() const {
   return restrict_to_whitelist_entries_;
 }
 
+void UkmService::SetInitializationCompleteCallbackForTesting(base::OnceClosure callback) {
+  if (initialize_complete_) {
+    std::move(callback).Run();
+  } else {
+    // Store the callback to be invoked when initialization is complete later.
+    initialization_complete_callback_ = std::move(callback);
+  }
+}
+
 }  // namespace ukm

components/ukm/ukm_service.h

@@ -121,6 +121,8 @@ class UkmService : public UkmRecorderImpl {
   // ukm::UkmRecorderImpl:
   bool ShouldRestrictToWhitelistedEntries() const override;
 
+  void SetInitializationCompleteCallbackForTesting(base::OnceClosure callback);
+
   // A weak pointer to the PrefService used to read and write preferences.
   PrefService* pref_service_;
 
@@ -154,6 +156,9 @@ class UkmService : public UkmRecorderImpl {
   bool initialize_started_ = false;
   bool initialize_complete_ = false;
 
+  // A callback invoked when initialization of the service is complete.
+  base::OnceClosure initialization_complete_callback_;
+
   SEQUENCE_CHECKER(sequence_checker_);
 
   // Weak pointers factory used to post task on different threads. All weak