Introduce ThreadTaskRunnerHandle::OverrideForTesting and TestMockTimeTaskRunner::ScopedContext.

base/test/test_mock_time_task_runner.cc

Index: base/test/test_mock_time_task_runner.cc
diff --git a/base/test/test_mock_time_task_runner.cc b/base/test/test_mock_time_task_runner.cc
index eaaceb6338166859c6da399b8de5edd30e665836..5d08976049a77de9fcb093b641aeba7aa67df3d6 100644
--- a/base/test/test_mock_time_task_runner.cc
+++ b/base/test/test_mock_time_task_runner.cc
@@ -8,6 +8,7 @@
 #include "base/macros.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/ref_counted.h"
+#include "base/threading/thread_task_runner_handle.h"
 #include "base/time/clock.h"
 #include "base/time/tick_clock.h"
 
@@ -120,6 +121,16 @@ TestMockTimeTaskRunner::TestOrderedPendingTask::operator=(
 
 // TestMockTimeTaskRunner -----------------------------------------------------
 
+// TODO(gab): This should also set the SequenceToken for the current thread.
+// Ref. TestMockTimeTaskRunner::RunsTasksOnCurrentThread().
+TestMockTimeTaskRunner::ScopedContext::ScopedContext(
+    scoped_refptr<TestMockTimeTaskRunner> scope)
+    : on_destroy_(ThreadTaskRunnerHandle::OverrideForTesting(scope)) {
+  scope->RunUntilIdle();
+}
+
+TestMockTimeTaskRunner::ScopedContext::~ScopedContext() = default;
+
 bool TestMockTimeTaskRunner::TemporalOrder::operator()(
     const TestOrderedPendingTask& first_task,
     const TestOrderedPendingTask& second_task) const {
@@ -213,6 +224,9 @@ TimeDelta TestMockTimeTaskRunner::NextPendingTaskDelay() const {
                         : tasks_.top().GetTimeToRun() - now_ticks_;
 }
 
+// TODO(gab): Combine |thread_checker_| with a SequenceToken to differentiate
+// between tasks running in the scope of this TestMockTimeTaskRunner and other
+// task runners sharing this thread. http://crbug.com/631186
 bool TestMockTimeTaskRunner::RunsTasksOnCurrentThread() const {
   return thread_checker_.CalledOnValidThread();
 }
@@ -253,6 +267,15 @@ void TestMockTimeTaskRunner::OnAfterTaskRun() {
 
 void TestMockTimeTaskRunner::ProcessAllTasksNoLaterThan(TimeDelta max_delta) {
   DCHECK_GE(max_delta, TimeDelta());
+
+  // Multiple test task runners can share the same thread for determinism in
+  // unit tests. Make sure this TestMockTimeTaskRunner's tasks run in its scope.
+  ScopedClosureRunner undo_override;
+  if (!ThreadTaskRunnerHandle::IsSet() ||
+      ThreadTaskRunnerHandle::Get() != this) {
+    undo_override = ThreadTaskRunnerHandle::OverrideForTesting(this);
+  }
+
   const TimeTicks original_now_ticks = now_ticks_;
   while (!IsElapsingStopped()) {
     OnBeforeSelectingTask();