[ScopedTaskEnvironment] Fix MAIN_THREAD_ONLY + MOCK_TIME

|task_tracker_| will be null in that mode and would crash prior to
this CL.

Precursor to https://chromium-review.googlesource.com/c/chromium/src/+/1747171

R=dcheng@chromium.org

Bug: 992483
Change-Id: I96a188573b1d291c9dcf7074ab3edd655c680bc5
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1747246
Auto-Submit: Gabriel Charette <gab@chromium.org>
Commit-Queue: Daniel Cheng <dcheng@chromium.org>
Reviewed-by: Daniel Cheng <dcheng@chromium.org>
Cr-Commit-Position: refs/heads/master@{#685897}

base/test/scoped_task_environment.cc

@@ -547,6 +547,11 @@ bool ScopedTaskEnvironment::MainThreadIsIdle() const {
 void ScopedTaskEnvironment::RunUntilIdle() {
   DCHECK_CALLED_ON_VALID_THREAD(main_thread_checker_);
 
+  if (threading_mode_ == ThreadingMode::MAIN_THREAD_ONLY) {
+    RunLoop(RunLoop::Type::kNestableTasksAllowed).RunUntilIdle();
+    return;
+  }
+
   // TODO(gab): This can be heavily simplified to essentially:
   //     bool HasMainThreadTasks() {
   //      if (message_loop_)
@@ -635,15 +640,15 @@ void ScopedTaskEnvironment::FastForwardBy(TimeDelta delta) {
   DCHECK(mock_time_domain_);
   DCHECK_GE(delta, TimeDelta());
 
-  const bool could_run_tasks = task_tracker_->AllowRunTasks();
+  const bool could_run_tasks = task_tracker_ && task_tracker_->AllowRunTasks();
 
   const TimeTicks fast_forward_until = mock_time_domain_->NowTicks() + delta;
   do {
     RunUntilIdle();
   } while (mock_time_domain_->FastForwardToNextTaskOrCap(fast_forward_until) !=
            MockTimeDomain::NextTaskSource::kNone);
 
-  if (!could_run_tasks)
+  if (task_tracker_ && !could_run_tasks)
     task_tracker_->DisallowRunTasks();
 }
 

base/test/scoped_task_environment_unittest.cc

@@ -20,6 +20,7 @@
 #include "base/task/sequence_manager/time_domain.h"
 #include "base/task/thread_pool/thread_pool.h"
 #include "base/test/bind_test_util.h"
+#include "base/test/gtest_util.h"
 #include "base/test/mock_callback.h"
 #include "base/test/mock_log.h"
 #include "base/test/test_timeouts.h"
@@ -311,12 +312,6 @@ TEST_F(ScopedTaskEnvironmentTest,
       ScopedTaskEnvironment::TimeSource::MOCK_TIME);
 }
 
-TEST_F(ScopedTaskEnvironmentTest, SingleThreadShouldNotInitializeThreadPool) {
-  ScopedTaskEnvironment scoped_task_environment(
-      ScopedTaskEnvironment::ThreadingMode::MAIN_THREAD_ONLY);
-  EXPECT_THAT(ThreadPoolInstance::Get(), IsNull());
-}
-
 // Verify that the right MessagePump is instantiated under each MainThreadType.
 // This avoids having to run all other ScopedTaskEnvironmentTests in every
 // MainThreadType which is redundant (message loop and message pump tests
@@ -1119,5 +1114,49 @@ TEST_F(ScopedTaskEnvironmentTest, TimeSourceMockTimeAlsoMocksNow) {
   EXPECT_EQ(Time::Now(), start_time + kDelay);
 }
 
+TEST_F(ScopedTaskEnvironmentTest, SingleThread) {
+  ScopedTaskEnvironment scoped_task_environment(
+      ScopedTaskEnvironment::ThreadingMode::MAIN_THREAD_ONLY);
+  EXPECT_THAT(ThreadPoolInstance::Get(), IsNull());
+
+  bool ran = false;
+  ThreadTaskRunnerHandle::Get()->PostTask(
+      FROM_HERE, base::BindLambdaForTesting([&]() { ran = true; }));
+  RunLoop().RunUntilIdle();
+  EXPECT_TRUE(ran);
+
+  EXPECT_DCHECK_DEATH(PostTask(FROM_HERE, {ThreadPool()}, DoNothing()));
+}
+
+// Verify that traits other than ThreadingMode can be applied to
+// ScopedTaskEnvironment.
+TEST_F(ScopedTaskEnvironmentTest, SingleThreadMockTime) {
+  ScopedTaskEnvironment scoped_task_environment(
+      ScopedTaskEnvironment::ThreadingMode::MAIN_THREAD_ONLY,
+      ScopedTaskEnvironment::TimeSource::MOCK_TIME);
+
+  const TimeTicks start_time = TimeTicks::Now();
+
+  constexpr TimeDelta kDelay = TimeDelta::FromSeconds(100);
+
+  int counter = 0;
+  ThreadTaskRunnerHandle::Get()->PostDelayedTask(
+      FROM_HERE, base::BindLambdaForTesting([&]() { counter += 1; }), kDelay);
+  ThreadTaskRunnerHandle::Get()->PostTask(
+      FROM_HERE, base::BindLambdaForTesting([&]() { counter += 2; }));
+
+  int expected_value = 0;
+  EXPECT_EQ(expected_value, counter);
+
+  scoped_task_environment.RunUntilIdle();
+  expected_value += 2;
+  EXPECT_EQ(expected_value, counter);
+
+  scoped_task_environment.FastForwardUntilNoTasksRemain();
+  expected_value += 1;
+  EXPECT_EQ(expected_value, counter);
+  EXPECT_EQ(TimeTicks::Now(), start_time + kDelay);
+}
+
 }  // namespace test
 }  // namespace base