Introduce ThreadTaskRunnerHandle::OverrideForTesting and TestMockTimeTaskRunner::ScopedContext.

base/threading/thread_task_runner_handle.cc

Index: base/threading/thread_task_runner_handle.cc
diff --git a/base/threading/thread_task_runner_handle.cc b/base/threading/thread_task_runner_handle.cc
index 190e18ffc68557c64f81f5b766a952c37d3c34a7..d71cabb135141dfabe36ef8c0842a2c3ef57358f 100644
--- a/base/threading/thread_task_runner_handle.cc
+++ b/base/threading/thread_task_runner_handle.cc
@@ -6,8 +6,10 @@
 
 #include <utility>
 
+#include "base/bind.h"
 #include "base/lazy_instance.h"
 #include "base/logging.h"
+#include "base/memory/ptr_util.h"
 #include "base/threading/sequenced_task_runner_handle.h"
 #include "base/threading/thread_local.h"
 
@@ -32,6 +34,50 @@ bool ThreadTaskRunnerHandle::IsSet() {
   return !!lazy_tls_ptr.Pointer()->Get();
 }
 
+// static
+ScopedClosureRunner ThreadTaskRunnerHandle::OverrideForTesting(
+    scoped_refptr<SingleThreadTaskRunner> overriding_task_runner) {
+  // OverrideForTesting() is not compatible with a SequencedTaskRunnerHandle
+  // being set (but SequencedTaskRunnerHandle::IsSet() includes
+  // ThreadTaskRunnerHandle::IsSet() so that's discounted as the only valid
+  // excuse for it to be true). Sadly this means that tests that merely need a
+  // SequencedTaskRunnerHandle on their main thread can be forced to use a
+  // ThreadTaskRunnerHandle if they're also using test task runners (that
+  // OverrideForTesting() when running their tasks from said main thread). To
+  // solve this: sequence_task_runner_handle.cc and thread_task_runner_handle.cc
+  // would have to be merged into a single impl file and share TLS state. This
+  // was deemed unecessary for now as most tests should use higher level
+  // constructs and not have to instantiate task runner handles on their own.
+  DCHECK(!SequencedTaskRunnerHandle::IsSet() || IsSet());
+
+  if (!IsSet()) {
+    std::unique_ptr<ThreadTaskRunnerHandle> top_level_ttrh =
+        MakeUnique<ThreadTaskRunnerHandle>(std::move(overriding_task_runner));
+    return ScopedClosureRunner(base::Bind(
+        [](std::unique_ptr<ThreadTaskRunnerHandle> ttrh_to_release) {},
+        base::Passed(&top_level_ttrh)));
+  }
+
+  ThreadTaskRunnerHandle* ttrh = lazy_tls_ptr.Pointer()->Get();
+  // Swap the two (and below bind |overriding_task_runner|, which is now the
+  // previous one, as the |task_runner_to_restore|).
+  ttrh->task_runner_.swap(overriding_task_runner);
+
+  return ScopedClosureRunner(base::Bind(
+      [](scoped_refptr<SingleThreadTaskRunner> task_runner_to_restore,
+         SingleThreadTaskRunner* expected_task_runner_before_restore) {
+        ThreadTaskRunnerHandle* ttrh = lazy_tls_ptr.Pointer()->Get();
+
+        DCHECK_EQ(expected_task_runner_before_restore, ttrh->task_runner_.get())
+            << "Nested overrides must expire their ScopedClosureRunners "
+               "in LIFO order.";
+
+        ttrh->task_runner_.swap(task_runner_to_restore);
+      },
+      base::Passed(&overriding_task_runner),
+      base::Unretained(ttrh->task_runner_.get())));
+}
+
 ThreadTaskRunnerHandle::ThreadTaskRunnerHandle(
     scoped_refptr<SingleThreadTaskRunner> task_runner)
     : task_runner_(std::move(task_runner)) {