Add support to base::Timer for custom task runners.

base/timer/timer_unittest.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 "base/memory/scoped_ptr.h"
 #include "base/message_loop/message_loop.h"
+#include "base/test/test_simple_task_runner.h"
 #include "base/timer/timer.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 using base::TimeDelta;
+using base::SingleThreadTaskRunner;
 
 namespace {
 
 // The message loops on which each timer should be tested.
 const base::MessageLoop::Type testing_message_loops[] = {
   base::MessageLoop::TYPE_DEFAULT,
   base::MessageLoop::TYPE_IO,
 #if !defined(OS_IOS)  // iOS does not allow direct running of the UI loop.
   base::MessageLoop::TYPE_UI,
 #endif
 };
 
 const int kNumTestingMessageLoops = arraysize(testing_message_loops);
 
 class OneShotTimerTester {
  public:
   explicit OneShotTimerTester(bool* did_run, unsigned milliseconds = 10)
       : did_run_(did_run),
-        delay_ms_(milliseconds) {
+        delay_ms_(milliseconds),
+        quit_message_loop_(true) {
   }
+
   void Start() {
     timer_.Start(FROM_HERE, TimeDelta::FromMilliseconds(delay_ms_), this,
                  &OneShotTimerTester::Run);
   }
+
+  void SetTaskRunner(scoped_refptr<SingleThreadTaskRunner> task_runner) {
+    quit_message_loop_ = false;
+    timer_.SetTaskRunner(task_runner);
+  }
+
  private:
   void Run() {
     *did_run_ = true;
-    base::MessageLoop::current()->QuitWhenIdle();
+    if (quit_message_loop_) {
+      base::MessageLoop::current()->QuitWhenIdle();
+    }
   }
+
   bool* did_run_;
   base::OneShotTimer<OneShotTimerTester> timer_;
   const unsigned delay_ms_;
+  bool quit_message_loop_;
 };
 
 class OneShotSelfDeletingTimerTester {
  public:
   explicit OneShotSelfDeletingTimerTester(bool* did_run) :
       did_run_(did_run),
       timer_(new base::OneShotTimer<OneShotSelfDeletingTimerTester>()) {
   }
+
   void Start() {
     timer_->Start(FROM_HERE, TimeDelta::FromMilliseconds(10), this,
                   &OneShotSelfDeletingTimerTester::Run);
   }
+
  private:
   void Run() {
     *did_run_ = true;
     timer_.reset();
     base::MessageLoop::current()->QuitWhenIdle();
   }
+
   bool* did_run_;
   scoped_ptr<base::OneShotTimer<OneShotSelfDeletingTimerTester> > timer_;
 };
 
 class RepeatingTimerTester {
  public:
   explicit RepeatingTimerTester(bool* did_run, const TimeDelta& delay)
       : did_run_(did_run), counter_(10), delay_(delay) {
   }
 
   void Start() {
     timer_.Start(FROM_HERE, delay_, this, &RepeatingTimerTester::Run);
   }
+
  private:
   void Run() {
     if (--counter_ == 0) {
       *did_run_ = true;
       timer_.Stop();
       base::MessageLoop::current()->QuitWhenIdle();
     }
   }
+
   bool* did_run_;
   int counter_;
   TimeDelta delay_;
   base::RepeatingTimer<RepeatingTimerTester> timer_;
 };
 
 void RunTest_OneShotTimer(base::MessageLoop::Type message_loop_type) {
   base::MessageLoop loop(message_loop_type);
 
   bool did_run = false;
@@ skipping 211 matching lines @@
 }
 
 // If underline timer does not handle properly, we will crash or fail
 // in full page heap environment.
 TEST(TimerTest, OneShotSelfDeletingTimer) {
   for (int i = 0; i < kNumTestingMessageLoops; i++) {
     RunTest_OneShotSelfDeletingTimer(testing_message_loops[i]);
   }
 }
 
+TEST(TimerTest, OneShotTimer_CustomTaskRunner) {
+  scoped_refptr<base::TestSimpleTaskRunner> task_runner =
+      new base::TestSimpleTaskRunner();
+
+  bool did_run = false;
+  OneShotTimerTester f(&did_run);
+  f.SetTaskRunner(task_runner);
+  f.Start();
+
+  EXPECT_FALSE(did_run);
+  task_runner->RunUntilIdle();
+  EXPECT_TRUE(did_run);
+}
+
 TEST(TimerTest, RepeatingTimer) {
   for (int i = 0; i < kNumTestingMessageLoops; i++) {
     RunTest_RepeatingTimer(testing_message_loops[i],
                            TimeDelta::FromMilliseconds(10));
   }
 }
 
 TEST(TimerTest, RepeatingTimer_Cancel) {
   for (int i = 0; i < kNumTestingMessageLoops; i++) {
     RunTest_RepeatingTimer_Cancel(testing_message_loops[i],
@@ skipping 176 matching lines @@
                 base::Bind(&SetCallbackHappened1));
     timer.Reset();
     // Since Reset happened before task ran, the user_task must not be cleared:
     ASSERT_FALSE(timer.user_task().is_null());
     base::MessageLoop::current()->Run();
     EXPECT_TRUE(g_callback_happened1);
   }
 }
 
 }  // namespace