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 = 2;
+
 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)
+  MessageLoop::RunLoop run_loop;
+  RunRunLoop(run_loop.AsWeakPtr());
+}
+
+void RunRunLoop(const base::WeakPtr<MessageLoop::RunLoop>& run_loop) {
+  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();
   }
 }
 
+static void DeferredQuitRunLoop(
+    const base::WeakPtr<MessageLoop::RunLoop>& run_loop,
+    int num_quit_deferrals) {
+  if (num_quit_deferrals <= 0) {
+    if (run_loop) run_loop->Stop();
+  } else {
+    MessageLoop::current()->PostTask(FROM_HERE,
+        base::Bind(&DeferredQuitRunLoop, run_loop, num_quit_deferrals - 1));
+  }
+}
+
+// Quit the MessageLoop Run of the specified |run_id|.
+void QuitRunLoop(const base::WeakPtr<MessageLoop::RunLoop>& run_loop) {
+  DeferredQuitRunLoop(run_loop, kNumQuitDeferrals);
+}
+
+base::Closure QuitRunLoopClosure(
+    const base::WeakPtr<MessageLoop::RunLoop>& run_loop) {
+  return base::Bind(&QuitRunLoop, run_loop);
+}
+
 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;
   }
+
+  MessageLoop::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.StopClosure()));
+  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);
+  MessageLoop::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);
+  MessageLoop::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) {
+    MessageLoop::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;
+  MessageLoop::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;
   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;
+  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