1 // Copyright 2018 The Chromium Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 #ifndef NET_TEST_TEST_WITH_TASK_ENVIRONMENT_H_ 6 #define NET_TEST_TEST_WITH_TASK_ENVIRONMENT_H_ 7 8 #include "base/test/task_environment.h" 9 #include "base/time/time.h" 10 #include "testing/gtest/include/gtest/gtest.h" 11 12 namespace base { 13 class TickClock; 14 } // namespace base 15 16 namespace net { 17 18 // Inherit from this class if a TaskEnvironment is needed in a test. 19 // Use in class hierachies where inheritance from ::testing::Test at the same 20 // time is not desirable or possible (for example, when inheriting from 21 // PlatformTest at the same time). 22 class WithTaskEnvironment { 23 public: 24 WithTaskEnvironment(const WithTaskEnvironment&) = delete; 25 WithTaskEnvironment& operator=(const WithTaskEnvironment&) = delete; 26 27 protected: 28 // Always uses MainThreadType::IO, |time_source| may optionally be provided 29 // to mock time. 30 explicit WithTaskEnvironment( 31 base::test::TaskEnvironment::TimeSource time_source = 32 base::test::TaskEnvironment::TimeSource::DEFAULT) task_environment_(base::test::TaskEnvironment::MainThreadType::IO,time_source)33 : task_environment_(base::test::TaskEnvironment::MainThreadType::IO, 34 time_source) {} 35 MainThreadIsIdle()36 [[nodiscard]] bool MainThreadIsIdle() const { 37 return task_environment_.MainThreadIsIdle(); 38 } 39 RunUntilIdle()40 void RunUntilIdle() { task_environment_.RunUntilIdle(); } 41 FastForwardBy(base::TimeDelta delta)42 void FastForwardBy(base::TimeDelta delta) { 43 task_environment_.FastForwardBy(delta); 44 } 45 FastForwardUntilNoTasksRemain()46 void FastForwardUntilNoTasksRemain() { 47 task_environment_.FastForwardUntilNoTasksRemain(); 48 } 49 50 // Only valid for instances using TimeSource::MOCK_TIME. AdvanceClock(base::TimeDelta delta)51 void AdvanceClock(base::TimeDelta delta) { 52 task_environment_.AdvanceClock(delta); 53 } 54 GetMockTickClock()55 [[nodiscard]] const base::TickClock* GetMockTickClock() { 56 return task_environment_.GetMockTickClock(); 57 } 58 GetPendingMainThreadTaskCount()59 [[nodiscard]] size_t GetPendingMainThreadTaskCount() const { 60 return task_environment_.GetPendingMainThreadTaskCount(); 61 } 62 NextMainThreadPendingTaskDelay()63 [[nodiscard]] base::TimeDelta NextMainThreadPendingTaskDelay() const { 64 return task_environment_.NextMainThreadPendingTaskDelay(); 65 } 66 67 private: 68 base::test::TaskEnvironment task_environment_; 69 }; 70 71 // Inherit from this class instead of ::testing::Test directly if a 72 // TaskEnvironment is needed in a test. 73 class TestWithTaskEnvironment : public ::testing::Test, 74 public WithTaskEnvironment { 75 public: 76 TestWithTaskEnvironment(const TestWithTaskEnvironment&) = delete; 77 TestWithTaskEnvironment& operator=(const TestWithTaskEnvironment&) = delete; 78 79 protected: 80 using WithTaskEnvironment::WithTaskEnvironment; 81 }; 82 83 } // namespace net 84 85 #endif // NET_TEST_TEST_WITH_TASK_ENVIRONMENT_H_ 86