libweave: Replace MessageLoop with weave::TaskRunner Removed dependency on base::MessageLoop and chromeos::MessageLoop. BUG=brillo:1257, brillo:1256 TEST=`FEATURES=test emerge-gizmo libweave buffet` Change-Id: Iff2f59c36016f930ed796bb5e0d83895fcb90013 Reviewed-on: https://chromium-review.googlesource.com/293420 Reviewed-by: Alex Vakulenko <avakulenko@chromium.org> Commit-Queue: Vitaly Buka <vitalybuka@chromium.org> Tested-by: Vitaly Buka <vitalybuka@chromium.org>

diff --git a/libweave/include/weave/mock_task_runner.h b/libweave/include/weave/mock_task_runner.h
new file mode 100644
index 0000000..313e05f
--- /dev/null
+++ b/libweave/include/weave/mock_task_runner.h

@@ -0,0 +1,82 @@
+// Copyright 2015 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef LIBWEAVE_INCLUDE_WEAVE_MOCK_TASK_RUNNER_H_
+#define LIBWEAVE_INCLUDE_WEAVE_MOCK_TASK_RUNNER_H_
+
+#include <weave/task_runner.h>
+
+#include <algorithm>
+#include <queue>
+#include <utility>
+#include <vector>
+
+#include <base/test/simple_test_clock.h>
+#include <gmock/gmock.h>
+
+namespace weave {
+namespace unittests {
+
+class MockTaskRunner : public TaskRunner {
+ public:
+  MockTaskRunner() {
+    test_clock_.SetNow(base::Time::Now());
+    using testing::_;
+    using testing::Invoke;
+    using testing::AnyNumber;
+    ON_CALL(*this, PostDelayedTask(_, _, _))
+        .WillByDefault(Invoke(this, &MockTaskRunner::SaveTask));
+    EXPECT_CALL(*this, PostDelayedTask(_, _, _)).Times(AnyNumber());
+  }
+  ~MockTaskRunner() override = default;
+
+  MOCK_METHOD3(PostDelayedTask,
+               void(const tracked_objects::Location&,
+                    const base::Closure&,
+                    base::TimeDelta));
+
+  bool RunOnce() {
+    if (queue_.empty())
+      return false;
+    auto top = queue_.top();
+    queue_.pop();
+    test_clock_.SetNow(std::max(test_clock_.Now(), top.first.first));
+    top.second.Run();
+    return true;
+  }
+
+  void Run() {
+    while (RunOnce()) {
+    }
+  }
+
+  base::SimpleTestClock* GetClock() { return &test_clock_; }
+
+ private:
+  void SaveTask(const tracked_objects::Location& from_here,
+                const base::Closure& task,
+                base::TimeDelta delay) {
+    queue_.emplace(std::make_pair(test_clock_.Now() + delay, ++counter_), task);
+  }
+
+  using QueueItem = std::pair<std::pair<base::Time, size_t>, base::Closure>;
+
+  struct Greater {
+    bool operator()(const QueueItem& a, const QueueItem& b) const {
+      return a.first > b.first;
+    }
+  };
+
+  size_t counter_{0};  // Keeps order of tasks with the same time.
+  base::SimpleTestClock test_clock_;
+
+  std::priority_queue<QueueItem,
+                      std::vector<QueueItem>,
+                      MockTaskRunner::Greater> queue_;
+};
+
+}  // namespace unittests
+}  // namespace weave
+
+#endif  // LIBWEAVE_INCLUDE_WEAVE_MOCK_TASK_RUNNER_H_

diff --git a/libweave/include/weave/task_runner.h b/libweave/include/weave/task_runner.h
index dcd76e4..2c3057c 100644
--- a/libweave/include/weave/task_runner.h
+++ b/libweave/include/weave/task_runner.h

@@ -5,12 +5,24 @@
 #ifndef LIBWEAVE_INCLUDE_WEAVE_TASK_RUNNER_H_
 #define LIBWEAVE_INCLUDE_WEAVE_TASK_RUNNER_H_
 
-#include <base/message_loop/message_loop.h>
-#include <chromeos/message_loops/message_loop.h>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <base/callback.h>
+#include <chromeos/errors/error.h>
 
 namespace weave {
 
-using TaskRunner = chromeos::MessageLoop;
+class TaskRunner {
+ public:
+  virtual void PostDelayedTask(const tracked_objects::Location& from_here,
+                               const base::Closure& task,
+                               base::TimeDelta delay) = 0;
+
+ protected:
+  virtual ~TaskRunner() = default;
+};
 
 }  // namespace weave
 

diff --git a/libweave/src/commands/cloud_command_proxy.h b/libweave/src/commands/cloud_command_proxy.h
index 98a7c6c..2dcb697 100644
--- a/libweave/src/commands/cloud_command_proxy.h
+++ b/libweave/src/commands/cloud_command_proxy.h

@@ -14,7 +14,6 @@
 #include <base/memory/weak_ptr.h>
 #include <base/scoped_observer.h>
 #include <weave/command.h>
-#include <weave/task_runner.h>
 
 #include "libweave/src/backoff_entry.h"
 #include "libweave/src/commands/cloud_command_update_interface.h"
@@ -23,6 +22,7 @@
 namespace weave {
 
 class CommandInstance;
+class TaskRunner;
 
 // Command proxy which publishes command updates to the cloud.
 class CloudCommandProxy final : public Command::Observer {

diff --git a/libweave/src/commands/cloud_command_proxy_unittest.cc b/libweave/src/commands/cloud_command_proxy_unittest.cc
index d98776a..cd3116b 100644
--- a/libweave/src/commands/cloud_command_proxy_unittest.cc
+++ b/libweave/src/commands/cloud_command_proxy_unittest.cc

@@ -7,11 +7,9 @@
 #include <memory>
 #include <queue>
 
-#include <base/test/simple_test_clock.h>
-#include <chromeos/message_loops/fake_message_loop.h>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include <chromeos/message_loops/fake_message_loop.h>
+#include <weave/mock_task_runner.h>
 
 #include "libweave/src/commands/command_dictionary.h"
 #include "libweave/src/commands/command_instance.h"
@@ -46,17 +44,6 @@
                     const base::Closure&));
 };
 
-// Mock-like task runner that allow the tests to inspect the calls to
-// TaskRunner::PostDelayedTask and verify the delays.
-class TestTaskRunner : public chromeos::FakeMessageLoop {
- public:
-  using FakeMessageLoop::FakeMessageLoop;
-  MOCK_METHOD3(PostDelayedTask,
-               TaskId(const tracked_objects::Location&,
-                      const base::Closure&,
-                      base::TimeDelta));
-};
-
 // Test back-off entry that uses the test clock.
 class TestBackoffEntry : public BackoffEntry {
  public:
@@ -89,19 +76,6 @@
     EXPECT_CALL(state_change_queue_, GetLastStateChangeId())
         .WillRepeatedly(testing::ReturnPointee(&current_state_update_id_));
 
-    auto on_post_task = [this](const tracked_objects::Location& from_here,
-                               const base::Closure& task,
-                               base::TimeDelta delay) -> bool {
-      clock_.Advance(delay);
-      task_queue_.push(task);
-      return true;
-    };
-
-    ON_CALL(task_runner_, PostDelayedTask(_, _, _))
-        .WillByDefault(testing::Invoke(on_post_task));
-
-    clock_.SetNow(base::Time::Now());
-
     // Set up the command schema.
     auto json = CreateDictionaryValue(R"({
       'calc': {
@@ -146,7 +120,7 @@
     static const BackoffEntry::Policy policy{0,     1000, 2.0,  0.0,
                                              20000, -1,   false};
     std::unique_ptr<TestBackoffEntry> backoff{
-        new TestBackoffEntry{&policy, &clock_}};
+        new TestBackoffEntry{&policy, task_runner_.GetClock()}};
 
     // Finally construct the CloudCommandProxy we are going to test here.
     std::unique_ptr<CloudCommandProxy> proxy{new CloudCommandProxy{
@@ -162,8 +136,7 @@
   base::CallbackList<void(StateChangeQueueInterface::UpdateID)> callbacks_;
   testing::StrictMock<MockCloudCommandUpdateInterface> cloud_updater_;
   testing::StrictMock<MockStateChangeQueueInterface> state_change_queue_;
-  base::SimpleTestClock clock_;
-  TestTaskRunner task_runner_{&clock_};
+  testing::StrictMock<unittests::MockTaskRunner> task_runner_;
   std::queue<base::Closure> task_queue_;
   CommandDictionary command_dictionary_;
   std::unique_ptr<CommandInstance> command_instance_;
@@ -251,8 +224,7 @@
   })";
   EXPECT_CALL(cloud_updater_, UpdateCommand(kCmdID, MatchJson(expect2), _, _))
       .WillOnce(SaveArg<3>(&on_error));
-  task_queue_.back().Run();
-  task_queue_.pop();
+  task_runner_.RunOnce();
 
   // Now backoff should be 2 seconds.
   expected_delay = base::TimeDelta::FromSeconds(2);
@@ -263,8 +235,7 @@
   base::Closure on_success;
   EXPECT_CALL(cloud_updater_, UpdateCommand(kCmdID, MatchJson(expect2), _, _))
       .WillOnce(SaveArg<2>(&on_success));
-  task_queue_.back().Run();
-  task_queue_.pop();
+  task_runner_.RunOnce();
 
   // Pretend it succeeds this time.
   on_success.Run();

diff --git a/libweave/src/device_manager.cc b/libweave/src/device_manager.cc
index 85da5e3..efbac77 100644
--- a/libweave/src/device_manager.cc
+++ b/libweave/src/device_manager.cc

@@ -6,8 +6,6 @@
 
 #include <string>
 
-#include <base/message_loop/message_loop.h>
-
 #include "libweave/src/base_api_handler.h"
 #include "libweave/src/buffet_config.h"
 #include "libweave/src/commands/command_manager.h"

diff --git a/libweave/src/device_registration_info.h b/libweave/src/device_registration_info.h
index fefc27c..f162ee4 100644
--- a/libweave/src/device_registration_info.h
+++ b/libweave/src/device_registration_info.h

@@ -23,7 +23,6 @@
 #include <weave/cloud.h>
 #include <weave/config.h>
 #include <weave/http_client.h>
-#include <weave/task_runner.h>
 
 #include "libweave/src/backoff_entry.h"
 #include "libweave/src/buffet_config.h"
@@ -47,6 +46,7 @@
 
 class Network;
 class StateManager;
+class TaskRunner;
 
 extern const char kErrorDomainOAuth2[];
 extern const char kErrorDomainGCD[];

diff --git a/libweave/src/device_registration_info_unittest.cc b/libweave/src/device_registration_info_unittest.cc
index c6f30e8..711755d 100644
--- a/libweave/src/device_registration_info_unittest.cc
+++ b/libweave/src/device_registration_info_unittest.cc

@@ -6,8 +6,6 @@
 
 #include <base/json/json_reader.h>
 #include <base/json/json_writer.h>
-#include <base/message_loop/message_loop.h>
-#include <base/run_loop.h>
 #include <base/values.h>
 #include <chromeos/bind_lambda.h>
 #include <chromeos/key_value_store.h>
@@ -162,22 +160,14 @@
   }
 
   bool RefreshAccessToken(chromeos::ErrorPtr* error) const {
-    base::MessageLoopForIO message_loop;
-    base::RunLoop run_loop;
-
     bool succeeded = false;
-    auto on_success = [&run_loop, &succeeded]() {
-      succeeded = true;
-      run_loop.Quit();
-    };
-    auto on_failure = [&run_loop, &error](const chromeos::Error* in_error) {
+    auto on_success = [&succeeded]() { succeeded = true; };
+    auto on_failure = [&error](const chromeos::Error* in_error) {
       if (error)
         *error = in_error->Clone();
-      run_loop.Quit();
     };
     dev_reg_->RefreshAccessToken(base::Bind(on_success),
                                  base::Bind(on_failure));
-    run_loop.Run();
     return succeeded;
   }
 
@@ -334,24 +324,17 @@
         return ReplyWithJson(200, json);
       })));
 
-  base::MessageLoopForIO message_loop;
-  base::RunLoop run_loop;
-
   bool succeeded = false;
-  auto on_success = [&run_loop, &succeeded,
-                     this](const base::DictionaryValue& info) {
+  auto on_success = [&succeeded, this](const base::DictionaryValue& info) {
     std::string id;
     EXPECT_TRUE(info.GetString("id", &id));
     EXPECT_EQ(test_data::kDeviceId, id);
     succeeded = true;
-    run_loop.Quit();
   };
-  auto on_failure = [&run_loop](const chromeos::Error* error) {
-    run_loop.Quit();
+  auto on_failure = [](const chromeos::Error* error) {
     FAIL() << "Should not be called";
   };
   dev_reg_->GetDeviceInfo(base::Bind(on_success), base::Bind(on_failure));
-  run_loop.Run();
   EXPECT_TRUE(succeeded);
 }
 

diff --git a/libweave/src/notification/pull_channel.cc b/libweave/src/notification/pull_channel.cc
index 7907f4a..d358e95 100644
--- a/libweave/src/notification/pull_channel.cc
+++ b/libweave/src/notification/pull_channel.cc

@@ -5,6 +5,7 @@
 #include "libweave/src/notification/pull_channel.h"
 
 #include <base/bind.h>
+#include <weave/task_runner.h>
 
 #include "libweave/src/notification/notification_delegate.h"
 

diff --git a/libweave/src/notification/pull_channel.h b/libweave/src/notification/pull_channel.h
index daa5832..7b10152 100644
--- a/libweave/src/notification/pull_channel.h
+++ b/libweave/src/notification/pull_channel.h

@@ -10,14 +10,13 @@
 
 #include <base/macros.h>
 #include <base/memory/weak_ptr.h>
-#include <base/single_thread_task_runner.h>
-#include <base/timer/timer.h>
-#include <weave/task_runner.h>
 
 #include "libweave/src/notification/notification_channel.h"
 
 namespace weave {
 
+class TaskRunner;
+
 class PullChannel : public NotificationChannel {
  public:
   PullChannel(base::TimeDelta pull_interval, TaskRunner* task_runner);

diff --git a/libweave/src/notification/xmpp_channel.cc b/libweave/src/notification/xmpp_channel.cc
index 7d0d5c6..13a3565 100644
--- a/libweave/src/notification/xmpp_channel.cc
+++ b/libweave/src/notification/xmpp_channel.cc

@@ -126,9 +126,9 @@
     // However, if the connection has never been established yet (e.g.
     // authorization failed), do not restart right now. Wait till we get
     // new credentials.
-    task_runner_->PostTask(
+    task_runner_->PostDelayedTask(
         FROM_HERE,
-        base::Bind(&XmppChannel::Restart, task_ptr_factory_.GetWeakPtr()));
+        base::Bind(&XmppChannel::Restart, task_ptr_factory_.GetWeakPtr()), {});
   } else if (delegate_) {
     delegate_->OnPermanentFailure();
   }
@@ -139,10 +139,11 @@
   // from expat XML parser and some stanza could cause the XMPP stream to be
   // reset and the parser to be re-initialized. We don't want to destroy the
   // parser while it is performing a callback invocation.
-  task_runner_->PostTask(
+  task_runner_->PostDelayedTask(
       FROM_HERE,
       base::Bind(&XmppChannel::HandleStanza, task_ptr_factory_.GetWeakPtr(),
-                 base::Passed(std::move(stanza))));
+                 base::Passed(std::move(stanza))),
+      {});
 }
 
 void XmppChannel::HandleStanza(std::unique_ptr<XmlNode> stanza) {

diff --git a/libweave/src/notification/xmpp_channel.h b/libweave/src/notification/xmpp_channel.h
index d6ae529..2f26748 100644
--- a/libweave/src/notification/xmpp_channel.h
+++ b/libweave/src/notification/xmpp_channel.h

@@ -23,6 +23,7 @@
 namespace weave {
 
 class Network;
+class TaskRunner;
 
 // Simple interface to abstract XmppChannel's SendMessage() method.
 class XmppChannelInterface {

diff --git a/libweave/src/notification/xmpp_channel_unittest.cc b/libweave/src/notification/xmpp_channel_unittest.cc
index d64c93a..173600e 100644
--- a/libweave/src/notification/xmpp_channel_unittest.cc
+++ b/libweave/src/notification/xmpp_channel_unittest.cc

@@ -9,8 +9,10 @@
 
 #include <base/test/simple_test_clock.h>
 #include <chromeos/bind_lambda.h>
-#include <chromeos/message_loops/fake_message_loop.h>
 #include <gtest/gtest.h>
+#include <weave/mock_task_runner.h>
+
+using testing::StrictMock;
 
 namespace weave {
 
@@ -157,11 +159,6 @@
 
 class XmppChannelTest : public ::testing::Test {
  protected:
-  void SetUp() override {
-    fake_loop_.SetAsCurrent();
-    clock_.SetNow(base::Time::Now());
-  }
-
   void StartStream() {
     xmpp_client_.fake_stream_.ExpectWritePacketString({}, kStartStreamMessage);
     xmpp_client_.fake_stream_.AddReadPacketString({}, kStartStreamResponse);
@@ -177,13 +174,12 @@
 
   void RunUntil(XmppChannel::XmppState st) {
     for (size_t n = 15; n && xmpp_client_.state() != st; --n)
-      fake_loop_.RunOnce(true);
+      task_runner_.RunOnce();
     EXPECT_EQ(st, xmpp_client_.state());
   }
 
-  base::SimpleTestClock clock_;
-  chromeos::FakeMessageLoop fake_loop_{&clock_};
-  FakeXmppChannel xmpp_client_{&fake_loop_};
+  StrictMock<unittests::MockTaskRunner> task_runner_;
+  FakeXmppChannel xmpp_client_{&task_runner_};
 };
 
 TEST_F(XmppChannelTest, StartStream) {

diff --git a/libweave/src/notification/xmpp_iq_stanza_handler.cc b/libweave/src/notification/xmpp_iq_stanza_handler.cc
index a61adc7..8803aae 100644
--- a/libweave/src/notification/xmpp_iq_stanza_handler.cc
+++ b/libweave/src/notification/xmpp_iq_stanza_handler.cc

@@ -110,8 +110,9 @@
     }
     auto p = requests_.find(id);
     if (p != requests_.end()) {
-      task_runner_->PostTask(
-          FROM_HERE, base::Bind(p->second, base::Passed(std::move(stanza))));
+      task_runner_->PostDelayedTask(
+          FROM_HERE, base::Bind(p->second, base::Passed(std::move(stanza))),
+          {});
       requests_.erase(p);
     }
   } else {

diff --git a/libweave/src/notification/xmpp_iq_stanza_handler.h b/libweave/src/notification/xmpp_iq_stanza_handler.h
index 3b3d71c..eb1f1b6 100644
--- a/libweave/src/notification/xmpp_iq_stanza_handler.h
+++ b/libweave/src/notification/xmpp_iq_stanza_handler.h

@@ -13,12 +13,12 @@
 #include <base/macros.h>
 #include <base/memory/weak_ptr.h>
 #include <base/single_thread_task_runner.h>
-#include <weave/task_runner.h>
 
 #include "libweave/src/notification/xmpp_stream_parser.h"
 
 namespace weave {
 
+class TaskRunner;
 class XmppChannelInterface;
 
 class IqStanzaHandler {

diff --git a/libweave/src/notification/xmpp_iq_stanza_handler_unittest.cc b/libweave/src/notification/xmpp_iq_stanza_handler_unittest.cc
index fd374b4..2078a6a 100644
--- a/libweave/src/notification/xmpp_iq_stanza_handler_unittest.cc
+++ b/libweave/src/notification/xmpp_iq_stanza_handler_unittest.cc

@@ -8,9 +8,9 @@
 #include <memory>
 
 #include <chromeos/bind_lambda.h>
-#include <chromeos/message_loops/mock_message_loop.h>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include <weave/mock_task_runner.h>
 
 #include "libweave/src/notification/xml_node.h"
 #include "libweave/src/notification/xmpp_channel.h"
@@ -75,14 +75,9 @@
 
 class IqStanzaHandlerTest : public testing::Test {
  public:
-  void SetUp() override {
-    mock_loop_.SetAsCurrent();
-  }
-
   testing::StrictMock<MockXmppChannelInterface> mock_xmpp_channel_;
-  base::SimpleTestClock clock_;
-  testing::NiceMock<chromeos::MockMessageLoop> mock_loop_{&clock_};
-  IqStanzaHandler iq_stanza_handler_{&mock_xmpp_channel_, &mock_loop_};
+  unittests::MockTaskRunner task_runner_;
+  IqStanzaHandler iq_stanza_handler_{&mock_xmpp_channel_, &task_runner_};
   MockResponseReceiver receiver_;
 };
 
@@ -106,7 +101,6 @@
   expected_msg = "<iq id='5' type='query' from='foo@bar' to='baz'><body/></iq>";
   EXPECT_CALL(mock_xmpp_channel_, SendMessage(expected_msg)).Times(1);
   iq_stanza_handler_.SendRequest("query", "foo@bar", "baz", "<body/>", {}, {});
-
   // This test ignores all the posted callbacks.
 }
 
@@ -128,8 +122,7 @@
 }
 
 TEST_F(IqStanzaHandlerTest, SequentialResponses) {
-  EXPECT_CALL(mock_loop_, PostDelayedTask(_, _, _))
-      .Times(2);
+  EXPECT_CALL(task_runner_, PostDelayedTask(_, _, _)).Times(2);
 
   EXPECT_CALL(mock_xmpp_channel_, SendMessage(_)).Times(2);
   iq_stanza_handler_.SendRequest("set", "", "", "<body/>",
@@ -137,8 +130,7 @@
   iq_stanza_handler_.SendRequest("get", "", "", "<body/>",
                                  receiver_.callback(2), {});
 
-  EXPECT_CALL(mock_loop_, PostDelayedTask(_, _, _))
-      .Times(2);
+  EXPECT_CALL(task_runner_, PostDelayedTask(_, _, _)).Times(2);
 
   EXPECT_CALL(receiver_, OnResponse(1, "foo"));
   auto request = XmlParser{}.Parse("<iq id='1' type='result'><foo/></iq>");
@@ -148,12 +140,11 @@
   request = XmlParser{}.Parse("<iq id='2' type='result'><bar/></iq>");
   EXPECT_TRUE(iq_stanza_handler_.HandleIqStanza(std::move(request)));
 
-  mock_loop_.Run();
+  task_runner_.Run();
 }
 
 TEST_F(IqStanzaHandlerTest, OutOfOrderResponses) {
-  EXPECT_CALL(mock_loop_, PostDelayedTask(_, _, _))
-      .Times(2);
+  EXPECT_CALL(task_runner_, PostDelayedTask(_, _, _)).Times(2);
 
   EXPECT_CALL(mock_xmpp_channel_, SendMessage(_)).Times(2);
   iq_stanza_handler_.SendRequest("set", "", "", "<body/>",
@@ -161,8 +152,7 @@
   iq_stanza_handler_.SendRequest("get", "", "", "<body/>",
                                  receiver_.callback(2), {});
 
-  EXPECT_CALL(mock_loop_, PostDelayedTask(_, _, _))
-      .Times(2);
+  EXPECT_CALL(task_runner_, PostDelayedTask(_, _, _)).Times(2);
 
   EXPECT_CALL(receiver_, OnResponse(2, "bar"));
   auto request = XmlParser{}.Parse("<iq id='2' type='result'><bar/></iq>");
@@ -172,12 +162,11 @@
   request = XmlParser{}.Parse("<iq id='1' type='result'><foo/></iq>");
   EXPECT_TRUE(iq_stanza_handler_.HandleIqStanza(std::move(request)));
 
-  mock_loop_.Run();
+  task_runner_.Run();
 }
 
 TEST_F(IqStanzaHandlerTest, RequestTimeout) {
-  EXPECT_CALL(mock_loop_, PostDelayedTask(_, _, _))
-      .Times(1);
+  EXPECT_CALL(task_runner_, PostDelayedTask(_, _, _)).Times(1);
 
   bool called = false;
   auto on_timeout = [&called]() { called = true; };
@@ -186,7 +175,7 @@
   EXPECT_FALSE(called);
   iq_stanza_handler_.SendRequest("set", "", "", "<body/>", {},
                                  base::Bind(on_timeout));
-  mock_loop_.Run();
+  task_runner_.Run();
   EXPECT_TRUE(called);
 }
 

diff --git a/libweave/src/privet/cloud_delegate.cc b/libweave/src/privet/cloud_delegate.cc
index f916187..f2cbdbf 100644
--- a/libweave/src/privet/cloud_delegate.cc
+++ b/libweave/src/privet/cloud_delegate.cc

@@ -10,7 +10,6 @@
 #include <base/bind.h>
 #include <base/logging.h>
 #include <base/memory/weak_ptr.h>
-#include <base/message_loop/message_loop.h>
 #include <base/values.h>
 #include <chromeos/errors/error.h>
 #include <weave/task_runner.h>

diff --git a/libweave/src/privet/cloud_delegate.h b/libweave/src/privet/cloud_delegate.h
index e31ba01..3b80b2e 100644
--- a/libweave/src/privet/cloud_delegate.h
+++ b/libweave/src/privet/cloud_delegate.h

@@ -12,7 +12,6 @@
 #include <base/callback.h>
 #include <base/memory/ref_counted.h>
 #include <base/observer_list.h>
-#include <weave/task_runner.h>
 
 #include "libweave/src/privet/privet_types.h"
 #include "libweave/src/privet/security_delegate.h"
@@ -26,6 +25,7 @@
 class CommandManager;
 class DeviceRegistrationInfo;
 class StateManager;
+class TaskRunner;
 
 namespace privet {
 

diff --git a/libweave/src/privet/privet_handler_unittest.cc b/libweave/src/privet/privet_handler_unittest.cc
index d9877a1..240da55 100644
--- a/libweave/src/privet/privet_handler_unittest.cc
+++ b/libweave/src/privet/privet_handler_unittest.cc

@@ -11,7 +11,6 @@
 #include <base/bind.h>
 #include <base/json/json_reader.h>
 #include <base/json/json_writer.h>
-#include <base/run_loop.h>
 #include <base/strings/string_util.h>
 #include <base/values.h>
 #include <gmock/gmock.h>
@@ -136,7 +135,6 @@
     handler_->HandleRequest(api, auth_header_, input,
                             base::Bind(&PrivetHandlerTest::HandlerCallback,
                                        base::Unretained(this)));
-    base::RunLoop().RunUntilIdle();
     return output_;
   }
 
@@ -152,7 +150,6 @@
     base::DictionaryValue dictionary;
     handler_->HandleRequest(api, auth_header_, &dictionary,
                             base::Bind(&PrivetHandlerTest::HandlerNoFound));
-    base::RunLoop().RunUntilIdle();
   }
 
   void SetNoWifiAndGcd() {
@@ -192,7 +189,6 @@
     EXPECT_EQ(404, status);
   }
 
-  base::MessageLoop message_loop_;
   std::unique_ptr<PrivetHandler> handler_;
   base::DictionaryValue output_;
   ConnectionState gcd_disabled_state_{ConnectionState::kDisabled};

diff --git a/libweave/src/privet/security_manager.cc b/libweave/src/privet/security_manager.cc
index 3921a45..cc5f64a 100644
--- a/libweave/src/privet/security_manager.cc
+++ b/libweave/src/privet/security_manager.cc

@@ -12,7 +12,6 @@
 #include <base/bind.h>
 #include <base/guid.h>
 #include <base/logging.h>
-#include <base/message_loop/message_loop.h>
 #include <base/rand_util.h>
 #include <base/stl_util.h>
 #include <base/strings/string_number_conversions.h>
@@ -21,8 +20,9 @@
 #include <chromeos/data_encoding.h>
 #include <chromeos/key_value_store.h>
 #include <chromeos/strings/string_utils.h>
-#include "libweave/external/crypto/p224_spake.h"
+#include <weave/task_runner.h>
 
+#include "libweave/external/crypto/p224_spake.h"
 #include "libweave/src/privet/constants.h"
 #include "libweave/src/privet/openssl_utils.h"
 

diff --git a/libweave/src/privet/security_manager.h b/libweave/src/privet/security_manager.h
index 9876781..5afeac5 100644
--- a/libweave/src/privet/security_manager.h
+++ b/libweave/src/privet/security_manager.h

@@ -16,7 +16,6 @@
 #include <base/memory/weak_ptr.h>
 #include <chromeos/errors/error.h>
 #include <chromeos/secure_blob.h>
-#include <weave/task_runner.h>
 
 #include "libweave/src/privet/security_delegate.h"
 
@@ -25,6 +24,9 @@
 }  // namespace crypto
 
 namespace weave {
+
+class TaskRunner;
+
 namespace privet {
 
 class SecurityManager : public SecurityDelegate {

diff --git a/libweave/src/privet/security_manager_unittest.cc b/libweave/src/privet/security_manager_unittest.cc
index dc86c13..c2024c2 100644
--- a/libweave/src/privet/security_manager_unittest.cc
+++ b/libweave/src/privet/security_manager_unittest.cc

@@ -15,18 +15,17 @@
 #include <base/bind.h>
 #include <base/files/file_util.h>
 #include <base/logging.h>
-#include <base/message_loop/message_loop.h>
 #include <base/rand_util.h>
 #include <base/strings/string_number_conversions.h>
 #include <base/strings/string_util.h>
 #include <chromeos/data_encoding.h>
 #include <chromeos/key_value_store.h>
-#include <chromeos/message_loops/fake_message_loop.h>
 #include <chromeos/strings/string_utils.h>
+#include "libweave/external/crypto/p224_spake.h"
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include <weave/mock_task_runner.h>
 
-#include "libweave/external/crypto/p224_spake.h"
 #include "libweave/src/privet/openssl_utils.h"
 
 using testing::Eq;
@@ -119,8 +118,7 @@
 
   const base::Time time_ = base::Time::FromTimeT(1410000000);
   base::FilePath embedded_code_path_{GetTempFilePath()};
-  base::SimpleTestClock clock_;
-  chromeos::FakeMessageLoop task_runner_{&clock_};
+  unittests::MockTaskRunner task_runner_;
   SecurityManager security_{{PairingType::kEmbeddedCode},
                             embedded_code_path_,
                             &task_runner_,

diff --git a/libweave/src/privet/wifi_bootstrap_manager.cc b/libweave/src/privet/wifi_bootstrap_manager.cc
index f8401bc..a3d6843 100644
--- a/libweave/src/privet/wifi_bootstrap_manager.cc
+++ b/libweave/src/privet/wifi_bootstrap_manager.cc

@@ -6,11 +6,11 @@
 
 #include <base/logging.h>
 #include <base/memory/weak_ptr.h>
-#include <base/message_loop/message_loop.h>
 #include <chromeos/bind_lambda.h>
 #include <chromeos/key_value_store.h>
 #include <weave/enum_to_string.h>
 #include <weave/network.h>
+#include <weave/task_runner.h>
 
 #include "libweave/src/privet/constants.h"
 
@@ -148,9 +148,10 @@
   if (new_state != state_) {
     state_ = new_state;
     // Post with weak ptr to avoid notification after this object destroyed.
-    task_runner_->PostTask(
+    task_runner_->PostDelayedTask(
         FROM_HERE, base::Bind(&WifiBootstrapManager::NotifyStateListeners,
-                              lifetime_weak_factory_.GetWeakPtr(), new_state));
+                              lifetime_weak_factory_.GetWeakPtr(), new_state),
+        {});
   } else {
     VLOG(3) << "Not notifying listeners of state change, "
             << "because the states are the same.";

diff --git a/libweave/src/privet/wifi_bootstrap_manager.h b/libweave/src/privet/wifi_bootstrap_manager.h
index 8f47223..aa47839 100644
--- a/libweave/src/privet/wifi_bootstrap_manager.h
+++ b/libweave/src/privet/wifi_bootstrap_manager.h

@@ -14,7 +14,6 @@
 #include <base/macros.h>
 #include <base/memory/weak_ptr.h>
 #include <base/scoped_observer.h>
-#include <weave/task_runner.h>
 
 #include "libweave/src/privet/cloud_delegate.h"
 #include "libweave/src/privet/privet_types.h"