Introduce ThreadTaskRunnerHandle::OverrideForTesting and TestMockTimeTaskRunner::ScopedContext.

base/test/test_mock_time_task_runner.h

Index: base/test/test_mock_time_task_runner.h
diff --git a/base/test/test_mock_time_task_runner.h b/base/test/test_mock_time_task_runner.h
index bf945a0b16f983c1bb09b57e22966da2d2aa091d..3246f89c64bbd249606ec67100329d36845d3c2e 100644
--- a/base/test/test_mock_time_task_runner.h
+++ b/base/test/test_mock_time_task_runner.h
@@ -17,6 +17,7 @@
 #include "base/synchronization/lock.h"
 #include "base/test/test_pending_task.h"
 #include "base/threading/thread_checker_impl.h"
+#include "base/threading/thread_task_runner_handle.h"
 #include "base/time/time.h"
 
 namespace base {
@@ -47,6 +48,52 @@ class TickClock;
 // that it supports running delayed tasks in the correct temporal order.
 class TestMockTimeTaskRunner : public SingleThreadTaskRunner {
  public:
+  // Everything that is executed in the scope of a ScopedContext will behave as
+  // though it ran under |scope| (i.e. ThreadTaskRunnerHandle,
+  // RunsTasksOnCurrentThread, etc.). This allows to have the test body all in

[danakj, 2017/02/16 17:57:24]

allows having the test body

or

allows the test body to be

[gab, 2017/02/16 21:03:36]

Done.

+  // one block when multiple TestMockTimeTaskRunners share the main thread. For
+  // example:
+  //
+  //   class ExampleFixture {
+  //    protected:
+  //     DoBarOnFoo() {
+  //       DCHECK(foo_task_runner_->RunsOnCurrentThread());
+  //       EXPECT_EQ(foo_task_runner_, ThreadTaskRunnerHandle::Get());
+  //       DoBar();
+  //     }
+  //
+  //     // Mock main task runner.
+  //     base::MessageLoop message_loop_;
+  //     base::ScopedMockTimeMessageLoopTaskRunner main_task_runner_;
+  //
+  //     // Mock foo task runner.
+  //     scoped_refptr<TestMockTimeTaskRunner> foo_task_runner_ =
+  //         new TestMockTimeTaskRunner();
+  //   };
+  //
+  //   TEST_F(ExampleFixture, DoBarOnFoo) {
+  //     DoThingsOnMain();
+  //     {
+  //       TestMockTimeTaskRunner::ScopedContext(foo_task_runner_.get());
+  //       DoBarOnFoo();
+  //     }
+  //     DoMoreThingsOnMain();
+  //   }
+  //
+  class ScopedContext {
+   public:
+    // Note: |scope| is ran until idle as part of this constructor to ensure
+    // that anything which runs in the underlying scope runs after any already
+    // pending tasks (the contrary would break the SequencedTraskRunner
+    // contract).
+    explicit ScopedContext(scoped_refptr<TestMockTimeTaskRunner> scope);
+    ~ScopedContext();
+
+   private:
+    ThreadTaskRunnerHandle::NestedForTesting task_runner_handle_;
+    DISALLOW_COPY_AND_ASSIGN(ScopedContext);
+  };
+
   // Constructs an instance whose virtual time will start at the Unix epoch, and
   // whose time ticks will start at zero.
   TestMockTimeTaskRunner();