1 files changed, 24 insertions, 25 deletions

diff --git a/chrome/browser/extensions/extension_browsertest.cc b/chrome/browser/extensions/extension_browsertest.cc
index a6280a2..90cca84 100644
--- a/chrome/browser/extensions/extension_browsertest.cc
+++ b/chrome/browser/extensions/extension_browsertest.cc
@@ -157,9 +157,8 @@ bool ExtensionBrowserTest::InstallOrUpdateExtension(const std::string& id,
 void ExtensionBrowserTest::ReloadExtension(const std::string& extension_id) {
   ExtensionsService* service = browser()->profile()->GetExtensionsService();
   service->ReloadExtension(extension_id);
-  ui_test_utils::RegisterAndWait(this,
-                                 NotificationType::EXTENSION_PROCESS_CREATED,
-                                 NotificationService::AllSources());
+  ui_test_utils::RegisterAndWait(NotificationType::EXTENSION_PROCESS_CREATED,
+                                 this);
 }
 
 void ExtensionBrowserTest::UnloadExtension(const std::string& extension_id) {
@@ -183,26 +182,30 @@ void ExtensionBrowserTest::EnableExtension(const std::string& extension_id) {
 }
 
 bool ExtensionBrowserTest::WaitForPageActionCountChangeTo(int count) {
+  NotificationRegistrar registrar;
+  registrar.Add(this,
+      NotificationType::EXTENSION_PAGE_ACTION_COUNT_CHANGED,
+      NotificationService::AllSources());
+
+  target_page_action_count_ = count;
   LocationBarTesting* location_bar =
       browser()->window()->GetLocationBar()->GetLocationBarForTesting();
-  if (location_bar->PageActionCount() != count) {
-    target_page_action_count_ = count;
-    ui_test_utils::RegisterAndWait(this,
-        NotificationType::EXTENSION_PAGE_ACTION_COUNT_CHANGED,
-        NotificationService::AllSources());
-  }
+  if (location_bar->PageActionCount() != count)
+    ui_test_utils::RunMessageLoop();
   return location_bar->PageActionCount() == count;
 }
 
 bool ExtensionBrowserTest::WaitForPageActionVisibilityChangeTo(int count) {
+  NotificationRegistrar registrar;
+  registrar.Add(this,
+      NotificationType::EXTENSION_PAGE_ACTION_VISIBILITY_CHANGED,
+      NotificationService::AllSources());
+
+  target_visible_page_action_count_ = count;
   LocationBarTesting* location_bar =
       browser()->window()->GetLocationBar()->GetLocationBarForTesting();
-  if (location_bar->PageActionVisibleCount() != count) {
-    target_visible_page_action_count_ = count;
-    ui_test_utils::RegisterAndWait(this,
-        NotificationType::EXTENSION_PAGE_ACTION_VISIBILITY_CHANGED,
-        NotificationService::AllSources());
-  }
+  if (location_bar->PageActionVisibleCount() != count)
+    ui_test_utils::RunMessageLoop();
   return location_bar->PageActionVisibleCount() == count;
 }
 
@@ -234,22 +237,19 @@ bool ExtensionBrowserTest::WaitForExtensionHostsToLoad() {
 
 bool ExtensionBrowserTest::WaitForExtensionInstall() {
   int before = extension_installs_observed_;
-  ui_test_utils::RegisterAndWait(this, NotificationType::EXTENSION_INSTALLED,
-                                 NotificationService::AllSources());
+  ui_test_utils::RegisterAndWait(NotificationType::EXTENSION_INSTALLED, this);
   return extension_installs_observed_ == (before + 1);
 }
 
 bool ExtensionBrowserTest::WaitForExtensionInstallError() {
   int before = extension_installs_observed_;
-  ui_test_utils::RegisterAndWait(this,
-                                 NotificationType::EXTENSION_INSTALL_ERROR,
-                                 NotificationService::AllSources());
+  ui_test_utils::RegisterAndWait(NotificationType::EXTENSION_INSTALL_ERROR,
+                                 this);
   return extension_installs_observed_ == before;
 }
 
 void ExtensionBrowserTest::WaitForExtensionLoad() {
-  ui_test_utils::RegisterAndWait(this, NotificationType::EXTENSION_LOADED,
-                                 NotificationService::AllSources());
+  ui_test_utils::RegisterAndWait(NotificationType::EXTENSION_LOADED, this);
   WaitForExtensionHostsToLoad();
 }
 
@@ -261,9 +261,8 @@ bool ExtensionBrowserTest::WaitForExtensionCrash(
     // The extension is already unloaded, presumably due to a crash.
     return true;
   }
-  ui_test_utils::RegisterAndWait(this,
-                                 NotificationType::EXTENSION_PROCESS_TERMINATED,
-                                 NotificationService::AllSources());
+  ui_test_utils::RegisterAndWait(NotificationType::EXTENSION_PROCESS_TERMINATED,
+                                 this);
   return (service->GetExtensionById(extension_id, true) == NULL);
 }