Add base::RunLoop and update ui_test_utils to use it to reduce flakiness

chrome/test/base/ui_test_utils.cc

 // Copyright (c) 2012 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 "chrome/test/base/ui_test_utils.h"
 
 #if defined(OS_WIN)
 #include <windows.h>
 #endif
 
@@ skipping 73 matching lines @@
 using content::NavigationController;
 using content::NavigationEntry;
 using content::OpenURLParams;
 using content::RenderViewHost;
 using content::RenderWidgetHost;
 using content::Referrer;
 using content::WebContents;
 
 static const int kDefaultWsPort = 8880;
 
+// Number of times to repost a QuitNow task so that the MessageLoop finishes up
+// pending tasks and tasks posted by those pending tasks without risking the
+// potential hang behavior of MessageLoop::QuitWhenIdle.
+static const int kNumQuitDeferrals = 5;
+
 namespace ui_test_utils {
 
 namespace {
 
 class DOMOperationObserver : public content::NotificationObserver,
                              public content::WebContentsObserver {
  public:
   explicit DOMOperationObserver(RenderViewHost* render_view_host)
       : content::WebContentsObserver(
             WebContents::FromRenderViewHost(render_view_host)),
         did_respond_(false) {
     registrar_.Add(this, content::NOTIFICATION_DOM_OPERATION_RESPONSE,
                    content::Source<RenderViewHost>(render_view_host));
-    ui_test_utils::RunMessageLoop();
+    message_loop_runner_ = new MessageLoopRunner;
+    message_loop_runner_->Run();
   }
 
   virtual void Observe(int type,
                        const content::NotificationSource& source,
                        const content::NotificationDetails& details) OVERRIDE {
     DCHECK(type == content::NOTIFICATION_DOM_OPERATION_RESPONSE);
     content::Details<DomOperationNotificationDetails> dom_op_details(details);
     response_ = dom_op_details->json;
     did_respond_ = true;
-    MessageLoopForUI::current()->Quit();
+    message_loop_runner_->QuitNow();
   }
 
   // Overridden from content::WebContentsObserver:
   virtual void RenderViewGone(base::TerminationStatus status) OVERRIDE {
-    MessageLoopForUI::current()->Quit();
+    message_loop_runner_->QuitNow();
   }
 
   bool GetResponse(std::string* response) WARN_UNUSED_RESULT {
     *response = response_;
     return did_respond_;
   }
 
  private:
   content::NotificationRegistrar registrar_;
   std::string response_;
   bool did_respond_;
+  scoped_refptr<MessageLoopRunner> message_loop_runner_;
 
   DISALLOW_COPY_AND_ASSIGN(DOMOperationObserver);
 };
 
 class FindInPageNotificationObserver : public content::NotificationObserver {
  public:
   explicit FindInPageNotificationObserver(TabContents* parent_tab)
       : parent_tab_(parent_tab),
         active_match_ordinal_(-1),
         number_of_matches_(0) {
     current_find_request_id_ =
         parent_tab->find_tab_helper()->current_find_request_id();
     registrar_.Add(this, chrome::NOTIFICATION_FIND_RESULT_AVAILABLE,
                    content::Source<WebContents>(parent_tab_->web_contents()));
-    ui_test_utils::RunMessageLoop();
+    message_loop_runner_ = new MessageLoopRunner;
+    message_loop_runner_->Run();
   }
 
   int active_match_ordinal() const { return active_match_ordinal_; }
 
   int number_of_matches() const { return number_of_matches_; }
 
   virtual void Observe(int type, const content::NotificationSource& source,
                        const content::NotificationDetails& details) {
     if (type == chrome::NOTIFICATION_FIND_RESULT_AVAILABLE) {
       content::Details<FindNotificationDetails> find_details(details);
       if (find_details->request_id() == current_find_request_id_) {
         // We get multiple responses and one of those will contain the ordinal.
         // This message comes to us before the final update is sent.
         if (find_details->active_match_ordinal() > -1)
           active_match_ordinal_ = find_details->active_match_ordinal();
         if (find_details->final_update()) {
           number_of_matches_ = find_details->number_of_matches();
-          MessageLoopForUI::current()->Quit();
+          message_loop_runner_->QuitNow();
         } else {
           DVLOG(1) << "Ignoring, since we only care about the final message";
         }
       }
     } else {
       NOTREACHED();
     }
   }
 
  private:
   content::NotificationRegistrar registrar_;
   TabContents* parent_tab_;
   // We will at some point (before final update) be notified of the ordinal and
   // we need to preserve it so we can send it later.
   int active_match_ordinal_;
   int number_of_matches_;
   // The id of the current find request, obtained from WebContents. Allows us
   // to monitor when the search completes.
   int current_find_request_id_;
+  scoped_refptr<MessageLoopRunner> message_loop_runner_;
 
   DISALLOW_COPY_AND_ASSIGN(FindInPageNotificationObserver);
 };
 
 class InProcessJavaScriptExecutionController
     : public base::RefCounted<InProcessJavaScriptExecutionController>,
       public JavaScriptExecutionController {
  public:
   explicit InProcessJavaScriptExecutionController(
       RenderViewHost* render_view_host)
@@ skipping 58 matching lines @@
   base::JSONReader reader(base::JSON_ALLOW_TRAILING_COMMAS);
   result->reset(reader.ReadToValue(json));
   if (!result->get()) {
     DLOG(ERROR) << reader.GetErrorMessage();
     return false;
   }
 
   return true;
 }
 
-void RunAllPendingMessageAndSendQuit(content::BrowserThread::ID thread_id) {
-  MessageLoop::current()->PostTask(FROM_HERE, MessageLoop::QuitClosure());
+void RunAllPendingMessageAndSendQuit(content::BrowserThread::ID thread_id,
+                                     const base::Closure& quit_task) {
+  MessageLoop::current()->PostTask(FROM_HERE,
+                                   MessageLoop::QuitWhenIdleClosure());
   RunMessageLoop();
-  content::BrowserThread::PostTask(thread_id, FROM_HERE,
-                                   MessageLoop::QuitClosure());
+  content::BrowserThread::PostTask(thread_id, FROM_HERE, quit_task);
 }
 
 }  // namespace
 
 void RunMessageLoop() {
-  MessageLoop* loop = MessageLoop::current();
-  MessageLoopForUI* ui_loop =
-      content::BrowserThread::CurrentlyOn(content::BrowserThread::UI) ?
-          MessageLoopForUI::current() : NULL;
-  MessageLoop::ScopedNestableTaskAllower allow(loop);
-  if (ui_loop) {
-#if defined(USE_AURA)
-    ui_loop->Run();
-#elif defined(TOOLKIT_VIEWS)
+  base::RunLoop run_loop;
+  RunRunLoop(run_loop.AsWeakPtr());
+}
+
+void RunRunLoop(const base::WeakPtr<base::RunLoop>& run_loop) {

[jar (doing other things), 2012/06/23 02:43:19]

nit: I didn't like the RunRun... prefix.  Perhaps better would be:
RunThisRunLoop
or
RunWeakRunLoop
or
RunWeakPtrRunLoop
or
other??

In truth, I also really wished this was a method on this weak-ptr... and then it
would just be a Run() method.  Is it worth making a local helper class
containing the weak pointer? Is that too clumsy? 

[jbates, 2012/06/25 20:04:57]

Renamed to RunThisRunLoop. I tried the helper class, but it didn't seem to be
worth its code-weight.

+  CHECK(run_loop);
+  MessageLoop::ScopedNestableTaskAllower allow(MessageLoop::current());
+  if (content::BrowserThread::CurrentlyOn(content::BrowserThread::UI)) {
+#if !defined(USE_AURA) && defined(TOOLKIT_VIEWS)
     views::AcceleratorHandler handler;
-    ui_loop->RunWithDispatcher(&handler);
-#elif defined(OS_POSIX) && !defined(OS_MACOSX) && !defined(OS_ANDROID)
-    ui_loop->RunWithDispatcher(NULL);
-#else
-    ui_loop->Run();
+    run_loop->set_dispatcher(&handler);
 #endif
+    run_loop->Run();
   } else {
-    loop->Run();
+    run_loop->Run();

[jar (doing other things), 2012/06/23 02:43:19]

Line 296 and 298 can now be pulled out of the if blocks, and then in fact, line
291 can now go inside the #if.

[jbates, 2012/06/25 20:04:57]

Done.

   }
 }
 
+static void DeferredQuitRunLoop(
+    const base::WeakPtr<base::RunLoop>& run_loop,
+    int num_quit_deferrals) {
+  if (num_quit_deferrals <= 0) {
+    if (run_loop) run_loop->Quit();

[jar (doing other things), 2012/06/23 02:43:19]

Do we gain value by iterating and posting when we don't have a run_loop any
more, or can we just return once it is gone?  After all, when we're done posting
(via line 308) and running, we'll just give up and do nothing anyway.
Could we just:
if (!run_loop)
  return;
before line 305?

It seems that if the run loop is gone, we'll just be posting into a messageloop
that has exited, and building cruft tasks there.

Or can we assert that we are the only way for this loop to finally end, and the
run_loop is always still intact?

[jbates, 2012/06/25 20:04:57]

Done.

+  } else {
+    MessageLoop::current()->PostTask(FROM_HERE,
+        base::Bind(&DeferredQuitRunLoop, run_loop, num_quit_deferrals - 1));
+  }
+}
+
+// Quit the MessageLoop Run of the specified |run_id|.

[jar (doing other things), 2012/06/23 02:43:19]

nit: update comment "run_id"

[jbates, 2012/06/25 20:04:57]

Done.

+void QuitRunLoop(const base::WeakPtr<base::RunLoop>& run_loop) {

[jar (doing other things), 2012/06/23 02:43:19]

This name is not as jarring... but if you did encapsulate the WeakPtr to a
RunLoop, this would indeed just be a Quit method... or perhaps a DeferredQuit()
method.

[jbates, 2012/06/25 20:04:57]

Changed to QuitThisRunLoop to be consistent with RunThisRunLoop.

+  DeferredQuitRunLoop(run_loop, kNumQuitDeferrals);
+}
+
 void RunAllPendingInMessageLoop() {
-  MessageLoop::current()->PostTask(FROM_HERE, MessageLoop::QuitClosure());
+  MessageLoop::current()->PostTask(FROM_HERE,
+                                   MessageLoop::QuitWhenIdleClosure());
   ui_test_utils::RunMessageLoop();
 }
 
 void RunAllPendingInMessageLoop(content::BrowserThread::ID thread_id) {
   if (content::BrowserThread::CurrentlyOn(thread_id)) {
     RunAllPendingInMessageLoop();
     return;
   }
   content::BrowserThread::ID current_thread_id;
   if (!content::BrowserThread::GetCurrentThreadIdentifier(&current_thread_id)) {
     NOTREACHED();
     return;
   }
+
+  base::RunLoop run_loop;
   content::BrowserThread::PostTask(thread_id, FROM_HERE,
-      base::Bind(&RunAllPendingMessageAndSendQuit, current_thread_id));
-
-  ui_test_utils::RunMessageLoop();
+      base::Bind(&RunAllPendingMessageAndSendQuit, current_thread_id,
+                 run_loop.QuitClosure()));
+  ui_test_utils::RunRunLoop(run_loop.AsWeakPtr());
 }
 
 bool GetCurrentTabTitle(const Browser* browser, string16* title) {
   WebContents* web_contents = browser->GetActiveWebContents();
   if (!web_contents)
     return false;
   NavigationEntry* last_entry = web_contents->GetController().GetActiveEntry();
   if (!last_entry)
     return false;
   title->assign(last_entry->GetTitleForDisplay(""));
   return true;
 }
 
 void WaitForNavigations(NavigationController* controller,
                         int number_of_navigations) {
   content::TestNavigationObserver observer(
       content::Source<NavigationController>(controller), NULL,
       number_of_navigations);
+  base::RunLoop run_loop;
   observer.WaitForObservation(
-      base::Bind(&ui_test_utils::RunMessageLoop),
-      base::Bind(&MessageLoop::Quit,
-                 base::Unretained(MessageLoopForUI::current())));
+      base::Bind(&ui_test_utils::RunRunLoop, run_loop.AsWeakPtr()),
+      base::Bind(&ui_test_utils::QuitRunLoop, run_loop.AsWeakPtr()));
 }
 
 void WaitForNewTab(Browser* browser) {
   WindowedNotificationObserver observer(
       chrome::NOTIFICATION_TAB_ADDED,
       content::Source<content::WebContentsDelegate>(browser));
   observer.Wait();
 }
 
 void WaitForLoadStop(WebContents* tab) {
@@ skipping 29 matching lines @@
   Browser::OpenURLOffTheRecord(profile, url);
   Browser* browser = browser::FindTabbedBrowser(
       profile->GetOffTheRecordProfile(), false);
   WaitForNavigations(&browser->GetActiveWebContents()->GetController(), 1);
 }
 
 void NavigateToURL(browser::NavigateParams* params) {
   content::TestNavigationObserver observer(
       content::NotificationService::AllSources(), NULL, 1);
   browser::Navigate(params);
+  base::RunLoop run_loop;
   observer.WaitForObservation(
-      base::Bind(&ui_test_utils::RunMessageLoop),
-      base::Bind(&MessageLoop::Quit,
-                 base::Unretained(MessageLoopForUI::current())));
-
+      base::Bind(&ui_test_utils::RunRunLoop, run_loop.AsWeakPtr()),
+      base::Bind(&ui_test_utils::QuitRunLoop, run_loop.AsWeakPtr()));
 }
 
 void NavigateToURL(Browser* browser, const GURL& url) {
   NavigateToURLWithDisposition(browser, url, CURRENT_TAB,
                                BROWSER_TEST_WAIT_FOR_NAVIGATION);
 }
 
 // Navigates the specified tab (via |disposition|) of |browser| to |url|,
 // blocking until the |number_of_navigations| specified complete.
 // |disposition| indicates what tab the download occurs in, and
@@ skipping 49 matching lines @@
         << " Unable to wait for navigation to \"" << url.spec()
         << "\" because the new tab is not available yet";
     return;
   } else if ((disposition == CURRENT_TAB) ||
       (disposition == NEW_FOREGROUND_TAB) ||
       (disposition == SINGLETON_TAB)) {
     // The currently selected tab is the right one.
     web_contents = browser->GetActiveWebContents();
   }
   if (disposition == CURRENT_TAB) {
+    base::RunLoop run_loop;
     same_tab_observer.WaitForObservation(
-        base::Bind(&ui_test_utils::RunMessageLoop),
-        base::Bind(&MessageLoop::Quit,
-                   base::Unretained(MessageLoopForUI::current())));
+        base::Bind(&ui_test_utils::RunRunLoop, run_loop.AsWeakPtr()),
+        base::Bind(&ui_test_utils::QuitRunLoop, run_loop.AsWeakPtr()));
     return;
   } else if (web_contents) {
     NavigationController* controller = &web_contents->GetController();
     WaitForNavigations(controller, number_of_navigations);
     return;
   }
   EXPECT_TRUE(NULL != web_contents) << " Unable to wait for navigation to \""
                                     << url.spec() << "\""
                                     << " because we can't get the tab contents";
 }
@@ skipping 223 matching lines @@
                      int type,
                      const content::NotificationSource& source) {
   content::NotificationRegistrar registrar;
   registrar.Add(observer, type, source);
   RunMessageLoop();
 }
 
 void WaitForBookmarkModelToLoad(BookmarkModel* model) {
   if (model->IsLoaded())
     return;
-  BookmarkLoadObserver observer;
+  base::RunLoop run_loop;
+  BookmarkLoadObserver observer(run_loop.AsWeakPtr());
   model->AddObserver(&observer);
-  RunMessageLoop();
+  RunRunLoop(run_loop.AsWeakPtr());
   model->RemoveObserver(&observer);
   ASSERT_TRUE(model->IsLoaded());
 }
 
 void WaitForTemplateURLServiceToLoad(TemplateURLService* service) {
   if (service->loaded())
     return;
   service->Load();
   TemplateURLServiceTestUtil::BlockTillServiceProcessesRequests();
   ASSERT_TRUE(service->loaded());
@@ skipping 37 matching lines @@
 
 bool SendKeyPressSync(const Browser* browser,
                       ui::KeyboardCode key,
                       bool control,
                       bool shift,
                       bool alt,
                       bool command) {
   gfx::NativeWindow window = NULL;
   if (!GetNativeWindow(browser, &window))
     return false;
-
+  scoped_refptr<MessageLoopRunner> runner =
+      new MessageLoopRunner;

[jar (doing other things), 2012/06/23 02:43:19]

nit: no need to wrap.

[jbates, 2012/06/25 20:04:57]

Done.

   bool result;
   result = ui_controls::SendKeyPressNotifyWhenDone(
-      window, key, control, shift, alt, command, MessageLoop::QuitClosure());
+      window, key, control, shift, alt, command, runner->QuitNowClosure());
 #if defined(OS_WIN)
   if (!result && BringBrowserWindowToFront(browser)) {
     result = ui_controls::SendKeyPressNotifyWhenDone(
-        window, key, control, shift, alt, command, MessageLoop::QuitClosure());
+        window, key, control, shift, alt, command, runner->QuitNowClosure());
   }
 #endif
   if (!result) {
     LOG(ERROR) << "ui_controls::SendKeyPressNotifyWhenDone failed";
     return false;
   }
 
   // Run the message loop. It'll stop running when either the key was received
   // or the test timed out (in which case testing::Test::HasFatalFailure should
   // be set).
-  RunMessageLoop();
+  runner->Run();
   return !testing::Test::HasFatalFailure();
 }
 
 bool SendKeyPressAndWait(const Browser* browser,
                          ui::KeyboardCode key,
                          bool control,
                          bool shift,
                          bool alt,
                          bool command,
                          int type,
                          const content::NotificationSource& source) {
   WindowedNotificationObserver observer(type, source);
 
   if (!SendKeyPressSync(browser, key, control, shift, alt, command))
     return false;
 
   observer.Wait();
   return !testing::Test::HasFatalFailure();
 }
 
 bool SendMouseMoveSync(const gfx::Point& location) {
-  if (!ui_controls::SendMouseMoveNotifyWhenDone(location.x(), location.y(),
-                                                MessageLoop::QuitClosure())) {
+  scoped_refptr<MessageLoopRunner> runner =
+      new MessageLoopRunner;

[jar (doing other things), 2012/06/23 02:43:19]

nit: un-wrap (here and below in file, like line 845, 1276, etc.)

[jbates, 2012/06/25 20:04:57]

Done.

+  if (!ui_controls::SendMouseMoveNotifyWhenDone(
+          location.x(), location.y(), runner->QuitNowClosure())) {
     return false;
   }
-  RunMessageLoop();
+  runner->Run();
   return !testing::Test::HasFatalFailure();
 }
 
 bool SendMouseEventsSync(ui_controls::MouseButton type, int state) {
+  scoped_refptr<MessageLoopRunner> runner =
+      new MessageLoopRunner;
   if (!ui_controls::SendMouseEventsNotifyWhenDone(
-          type, state, MessageLoop::QuitClosure())) {
+          type, state, runner->QuitNowClosure())) {
     return false;
   }
-  RunMessageLoop();
+  runner->Run();
   return !testing::Test::HasFatalFailure();
 }
 
-TimedMessageLoopRunner::TimedMessageLoopRunner()
-    : loop_(new MessageLoopForUI()),
-      owned_(true),
-      quit_loop_invoked_(false) {
+MessageLoopRunner::MessageLoopRunner() {
 }
 
-TimedMessageLoopRunner::~TimedMessageLoopRunner() {
-  if (owned_)
-    delete loop_;
+MessageLoopRunner::~MessageLoopRunner() {
 }
 
-void TimedMessageLoopRunner::RunFor(int ms) {
-  QuitAfter(ms);
-  quit_loop_invoked_ = false;
-  loop_->Run();
+void MessageLoopRunner::Run() {
+  ui_test_utils::RunRunLoop(run_loop_.AsWeakPtr());
 }
 
-void TimedMessageLoopRunner::Quit() {
-  quit_loop_invoked_ = true;
-  loop_->PostTask(FROM_HERE, MessageLoop::QuitClosure());
+base::Closure MessageLoopRunner::QuitNowClosure() {
+  return base::Bind(&MessageLoopRunner::QuitNow, this);
 }
 
-void TimedMessageLoopRunner::QuitAfter(int ms) {
-  quit_loop_invoked_ = true;
-  loop_->PostDelayedTask(
-      FROM_HERE,
-      MessageLoop::QuitClosure(),
-      base::TimeDelta::FromMilliseconds(ms));
+void MessageLoopRunner::QuitNow() {
+  ui_test_utils::QuitRunLoop(run_loop_.AsWeakPtr());
 }
 
 TestWebSocketServer::TestWebSocketServer()
     : started_(false),
       port_(kDefaultWsPort),
       secure_(false) {
 #if defined(OS_POSIX)
   process_group_id_ = base::kNullProcessHandle;
 #endif
 }
@@ skipping 144 matching lines @@
   registrar_.Add(this, notification_type, source);
 }
 
 WindowedNotificationObserver::~WindowedNotificationObserver() {}
 
 void WindowedNotificationObserver::Wait() {
   if (seen_)
     return;
 
   running_ = true;
-  ui_test_utils::RunMessageLoop();
+  message_loop_runner_ = new MessageLoopRunner;
+  message_loop_runner_->Run();
   EXPECT_TRUE(seen_);
 }
 
 void WindowedNotificationObserver::Observe(
     int type,
     const content::NotificationSource& source,
     const content::NotificationDetails& details) {
   source_ = source;
   details_ = details;
   seen_ = true;
   if (!running_)
     return;
 
-  MessageLoopForUI::current()->Quit();
+  message_loop_runner_->QuitNow();
   running_ = false;
 }
 
 WindowedTabAddedNotificationObserver::WindowedTabAddedNotificationObserver(
     const content::NotificationSource& source)
     : WindowedNotificationObserver(chrome::NOTIFICATION_TAB_ADDED, source),
       added_tab_(NULL) {
 }
 
 void WindowedTabAddedNotificationObserver::Observe(
@@ skipping 30 matching lines @@
   expected_titles_.push_back(expected_title);
 }
 
 TitleWatcher::~TitleWatcher() {
 }
 
 const string16& TitleWatcher::WaitAndGetTitle() {
   if (expected_title_observed_)
     return observed_title_;
   quit_loop_on_observation_ = true;
-  ui_test_utils::RunMessageLoop();
+  message_loop_runner_ = new MessageLoopRunner;
+  message_loop_runner_->Run();
   return observed_title_;
 }
 
 void TitleWatcher::Observe(int type,
                            const content::NotificationSource& source,
                            const content::NotificationDetails& details) {
   if (type == content::NOTIFICATION_WEB_CONTENTS_TITLE_UPDATED) {
     WebContents* source_contents = content::Source<WebContents>(source).ptr();
     ASSERT_EQ(web_contents_, source_contents);
   } else if (type == content::NOTIFICATION_LOAD_STOP) {
     NavigationController* controller =
         content::Source<NavigationController>(source).ptr();
     ASSERT_EQ(&web_contents_->GetController(), controller);
   } else {
     FAIL() << "Unexpected notification received.";
   }
 
   std::vector<string16>::const_iterator it =
       std::find(expected_titles_.begin(),
                 expected_titles_.end(),
                 web_contents_->GetTitle());
   if (it == expected_titles_.end())
     return;
   observed_title_ = *it;
   expected_title_observed_ = true;
-  if (quit_loop_on_observation_)
-    MessageLoopForUI::current()->Quit();
+  if (quit_loop_on_observation_) {
+    // Only call Quit once, on first Observe:
+    quit_loop_on_observation_ = false;
+    message_loop_runner_->QuitNow();
+  }
 }
 
 BrowserAddedObserver::BrowserAddedObserver()
     : notification_observer_(
           chrome::NOTIFICATION_BROWSER_OPENED,
           content::NotificationService::AllSources()) {
   original_browsers_.insert(BrowserList::begin(), BrowserList::end());
 }
 
 BrowserAddedObserver::~BrowserAddedObserver() {
@@ skipping 14 matching lines @@
 DOMMessageQueue::~DOMMessageQueue() {}
 
 void DOMMessageQueue::Observe(int type,
                               const content::NotificationSource& source,
                               const content::NotificationDetails& details) {
   content::Details<DomOperationNotificationDetails> dom_op_details(details);
   content::Source<RenderViewHost> sender(source);
   message_queue_.push(dom_op_details->json);
   if (waiting_for_message_) {
     waiting_for_message_ = false;
-    MessageLoopForUI::current()->Quit();
+    message_loop_runner_->QuitNow();
   }
 }
 
 void DOMMessageQueue::ClearQueue() {
   message_queue_ = std::queue<std::string>();
 }
 
 bool DOMMessageQueue::WaitForMessage(std::string* message) {
   if (message_queue_.empty()) {
     waiting_for_message_ = true;
     // This will be quit when a new message comes in.
-    RunMessageLoop();
+    message_loop_runner_ = new MessageLoopRunner;
+    message_loop_runner_->Run();
   }
   // The queue should not be empty, unless we were quit because of a timeout.
   if (message_queue_.empty())
     return false;
   if (message)
     *message = message_queue_.front();
   return true;
 }
 
 // Coordinates taking snapshots of a |RenderWidget|.
 class SnapshotTaker {
  public:
   SnapshotTaker() : bitmap_(NULL) {}
 
   bool TakeRenderWidgetSnapshot(RenderWidgetHost* rwh,
                                 const gfx::Size& page_size,
                                 const gfx::Size& desired_size,
                                 SkBitmap* bitmap) WARN_UNUSED_RESULT {
     bitmap_ = bitmap;
     ThumbnailGenerator* generator = g_browser_process->GetThumbnailGenerator();
     generator->MonitorRenderer(rwh, true);
     snapshot_taken_ = false;
     generator->AskForSnapshot(
         rwh,
         base::Bind(&SnapshotTaker::OnSnapshotTaken, base::Unretained(this)),
         page_size,
         desired_size);
-    ui_test_utils::RunMessageLoop();
+    message_loop_runner_ = new MessageLoopRunner;
+    message_loop_runner_->Run();
     return snapshot_taken_;
   }
 
   bool TakeEntirePageSnapshot(RenderViewHost* rvh,
                               SkBitmap* bitmap) WARN_UNUSED_RESULT {
     const wchar_t* script =
         L"window.domAutomationController.send("
         L"    JSON.stringify([document.width, document.height]))";
     std::string json;
     if (!ui_test_utils::ExecuteJavaScriptAndExtractString(
@@ skipping 14 matching lines @@
     // Take the snapshot.
     gfx::Size page_size(width, height);
     return TakeRenderWidgetSnapshot(rvh, page_size, page_size, bitmap);
   }
 
  private:
   // Called when the ThumbnailGenerator has taken the snapshot.
   void OnSnapshotTaken(const SkBitmap& bitmap) {
     *bitmap_ = bitmap;
     snapshot_taken_ = true;
-    MessageLoop::current()->Quit();
+    message_loop_runner_->QuitNow();
   }
 
   SkBitmap* bitmap_;
   // Whether the snapshot was actually taken and received by this SnapshotTaker.
   // This will be false if the test times out.
   bool snapshot_taken_;
+  scoped_refptr<MessageLoopRunner> message_loop_runner_;
 
   DISALLOW_COPY_AND_ASSIGN(SnapshotTaker);
 };
 
 bool TakeRenderWidgetSnapshot(RenderWidgetHost* rwh,
                               const gfx::Size& page_size,
                               SkBitmap* bitmap) {
   DCHECK(bitmap);
   SnapshotTaker taker;
   return taker.TakeRenderWidgetSnapshot(rwh, page_size, page_size, bitmap);
 }
 
 bool TakeEntirePageSnapshot(RenderViewHost* rvh, SkBitmap* bitmap) {
   DCHECK(bitmap);
   SnapshotTaker taker;
   return taker.TakeEntirePageSnapshot(rvh, bitmap);
 }
 
 void OverrideGeolocation(double latitude, double longitude) {
   content::Geoposition position;
   position.latitude = latitude;
   position.longitude = longitude;
   position.altitude = 0.;
   position.accuracy = 0.;
   position.timestamp = base::Time::Now();
+  scoped_refptr<MessageLoopRunner> runner =
+      new MessageLoopRunner;
   content::OverrideLocationForTesting(position,
-                                      base::Bind(MessageLoop::QuitClosure()));
-  RunMessageLoop();
+                                      runner->QuitNowClosure());
+  runner->Run();
 }
 
 namespace internal {
 
 void ClickTask(ui_controls::MouseButton button,
                int state,
                const base::Closure& followup) {
   if (!followup.is_null())
     ui_controls::SendMouseEventsNotifyWhenDone(button, state, followup);
   else
     ui_controls::SendMouseEvents(button, state);
 }
 
 }  // namespace internal
 }  // namespace ui_test_utils