1 files changed, 7 insertions, 1 deletions
diff --git a/remoting/host/chromoting_host_unittest.cc b/remoting/host/chromoting_host_unittest.cc
index 0842832..129d922 100644
--- a/remoting/host/chromoting_host_unittest.cc
+++ b/remoting/host/chromoting_host_unittest.cc
@@ -78,7 +78,7 @@ class ChromotingHostTest : public testing::Test {
     ON_CALL(context_, network_message_loop())
         .WillByDefault(Return(message_loop_proxy_.get()));
     ON_CALL(context_, ui_message_loop())
-        .WillByDefault(Return(message_loop_proxy_.get()));
+        .WillByDefault(Return(&message_loop_));
     EXPECT_CALL(context_, main_message_loop())
         .Times(AnyNumber());
     EXPECT_CALL(context_, encode_message_loop())
@@ -88,6 +88,12 @@ class ChromotingHostTest : public testing::Test {
     EXPECT_CALL(context_, ui_message_loop())
         .Times(AnyNumber());
 
+    context_.SetUITaskPostFunction(base::Bind(
+        static_cast<void(MessageLoop::*)(
+            const tracked_objects::Location&,
+            const base::Closure&)>(&MessageLoop::PostTask),
+        base::Unretained(&message_loop_)));
+
     Capturer* capturer = new CapturerFake();
     event_executor_ = new MockEventExecutor();
     curtain_ = new MockCurtain();