FOR REFERENCE ONLY Task Scheduler COM Task Runner

base/task_scheduler/scheduler_single_thread_task_runner_manager_unittest.cc

 // Copyright 2017 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/task_scheduler/scheduler_single_thread_task_runner_manager.h"
 
 #include "base/bind.h"
 #include "base/memory/ptr_util.h"
 #include "base/synchronization/lock.h"
 #include "base/synchronization/waitable_event.h"
 #include "base/task_scheduler/delayed_task_manager.h"
 #include "base/task_scheduler/post_task.h"
 #include "base/task_scheduler/scheduler_worker_pool_params.h"
 #include "base/task_scheduler/task_tracker.h"
 #include "base/task_scheduler/task_traits.h"
 #include "base/test/test_timeouts.h"
 #include "base/threading/simple_thread.h"
 #include "base/threading/thread.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
+#if defined(OS_WIN)
+#include <objbase.h>
+#include <windows.h>
+
+#include "base/win/current_module.h"
+#endif
+
 namespace base {
 namespace internal {
 
 namespace {
 
 enum WorkerPoolType : size_t {
   BACKGROUND_WORKER_POOL = 0,
   FOREGROUND_WORKER_POOL,
 };
 
@@ skipping 320 matching lines @@
 
   task_running.Wait();
   CallJoinFromDifferentThread join_from_different_thread(
       single_thread_task_runner_manager_.get());
   join_from_different_thread.Start();
   join_from_different_thread.WaitForRunToStart();
   task_blocking.Signal();
   join_from_different_thread.Join();
 }
 
+#if defined(OS_WIN)
+
+TEST_F(TaskSchedulerSingleThreadTaskRunnerManagerTest, COMSTAInitialized) {
+  scoped_refptr<SingleThreadTaskRunner> com_task_runner =
+      single_thread_task_runner_manager_->CreateCOMSTATaskRunnerWithTraits(
+          TaskTraits().WithShutdownBehavior(
+              TaskShutdownBehavior::BLOCK_SHUTDOWN));
+
+  com_task_runner->PostTask(
+      FROM_HERE, Bind([]() {
+        HRESULT hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
+        if (SUCCEEDED(hr)) {
+          ADD_FAILURE() << "COM STA was not initialized on this thread";
+          CoUninitialize();
+        }
+      }));
+
+  task_tracker_.Shutdown();
+}
+
+namespace {
+
+const wchar_t* const kTestWindowClassName =
+    L"TaskSchedulerSingleThreadTaskRunnerManagerTestWinMessageWindow";
+
+class TaskSchedulerSingleThreadTaskRunnerManagerTestWin
+    : public TaskSchedulerSingleThreadTaskRunnerManagerTest {
+ public:
+  TaskSchedulerSingleThreadTaskRunnerManagerTestWin() = default;
+
+  void SetUp() override {
+    TaskSchedulerSingleThreadTaskRunnerManagerTest::SetUp();
+    register_class_succeeded_ = RegisterTestWindowClass();
+    ASSERT_TRUE(register_class_succeeded_);
+  }
+
+  void TearDown() override {
+    if (register_class_succeeded_)
+      ::UnregisterClass(kTestWindowClassName, CURRENT_MODULE());
+
+    TaskSchedulerSingleThreadTaskRunnerManagerTest::TearDown();
+  }
+
+  HWND CreateTestWindow() {
+    return CreateWindow(kTestWindowClassName, kTestWindowClassName, 0, 0, 0, 0,
+                        0, HWND_MESSAGE, nullptr, CURRENT_MODULE(), nullptr);
+  }
+
+ private:
+  bool RegisterTestWindowClass() {
+    WNDCLASSEX window_class = {};
+    window_class.cbSize = sizeof(window_class);
+    window_class.lpfnWndProc = &::DefWindowProc;
+    window_class.hInstance = CURRENT_MODULE();
+    window_class.lpszClassName = kTestWindowClassName;
+    return !!::RegisterClassEx(&window_class);
+  }
+
+  bool register_class_succeeded_ = false;
+
+  DISALLOW_COPY_AND_ASSIGN(TaskSchedulerSingleThreadTaskRunnerManagerTestWin);
+};
+
+}  // namespace
+
+TEST_F(TaskSchedulerSingleThreadTaskRunnerManagerTestWin, PumpsMessages) {
+  scoped_refptr<SingleThreadTaskRunner> com_task_runner =
+      single_thread_task_runner_manager_->CreateCOMSTATaskRunnerWithTraits(
+          TaskTraits().WithShutdownBehavior(
+              TaskShutdownBehavior::BLOCK_SHUTDOWN));
+  HWND hwnd = nullptr;
+  // HWNDs process messages on the thread that created them, so we have to
+  // create them within the context of the task runner to properly simulate a
+  // COM callback.
+  com_task_runner->PostTask(
+      FROM_HERE,
+      Bind([](TaskSchedulerSingleThreadTaskRunnerManagerTestWin* test_harness,
+              HWND* hwnd) { *hwnd = test_harness->CreateTestWindow(); },
+           Unretained(this), &hwnd));
+
+  WaitableEvent wait_for_create_window(
+      WaitableEvent::ResetPolicy::MANUAL,
+      WaitableEvent::InitialState::NOT_SIGNALED);
+  com_task_runner->PostTask(
+      FROM_HERE,
+      Bind(&WaitableEvent::Signal, Unretained(&wait_for_create_window)));
+  wait_for_create_window.Wait();
+
+  ASSERT_NE(hwnd, nullptr);
+  // If the message pump isn't running, we will hang here. This simulates how
+  // COM would receive a callback with its own message HWND.
+  SendMessage(hwnd, WM_USER, 0, 0);
+
+  com_task_runner->PostTask(
+      FROM_HERE, Bind([](HWND hwnd) { ::DestroyWindow(hwnd); }, hwnd));
+
+  task_tracker_.Shutdown();
+}
+
+#endif  // defined(OS_WIN)
+
 }  // namespace internal
 }  // namespace base