diff --git a/buffet/base_api_handler_unittest.cc b/buffet/base_api_handler_unittest.cc
index aa780be..c7df9bb 100644
--- a/buffet/base_api_handler_unittest.cc
+++ b/buffet/base_api_handler_unittest.cc
@@ -53,7 +53,7 @@
         command_manager_, state_manager_,
         std::unique_ptr<BuffetConfig>{new BuffetConfig{
             std::unique_ptr<StorageInterface>{new MemStorage}}},
-        transport_, true, nullptr));
+        transport_, nullptr, true, nullptr));
     handler_.reset(new BaseApiHandler{
         dev_reg_->AsWeakPtr(), state_manager_, command_manager_});
   }
diff --git a/buffet/device_registration_info.cc b/buffet/device_registration_info.cc
index 4b74d7a..74e4433 100644
--- a/buffet/device_registration_info.cc
+++ b/buffet/device_registration_info.cc
@@ -11,7 +11,6 @@
 
 #include <base/bind.h>
 #include <base/json/json_writer.h>
-#include <base/message_loop/message_loop.h>
 #include <base/values.h>
 #include <chromeos/bind_lambda.h>
 #include <chromeos/data_encoding.h>
@@ -114,9 +113,11 @@
     const std::shared_ptr<StateManager>& state_manager,
     std::unique_ptr<BuffetConfig> config,
     const std::shared_ptr<chromeos::http::Transport>& transport,
+    const scoped_refptr<base::SingleThreadTaskRunner>& task_runner,
     bool notifications_enabled,
     privetd::ShillClient* shill_client)
     : transport_{transport},
+      task_runner_{task_runner},
       command_manager_{command_manager},
       state_manager_{state_manager},
       config_{std::move(config)},
@@ -183,8 +184,7 @@
 
 void DeviceRegistrationInfo::ScheduleStartDevice(const base::TimeDelta& later) {
   SetRegistrationStatus(RegistrationStatus::kConnecting);
-  base::MessageLoop* current = base::MessageLoop::current();
-  if (!current)
+  if (!task_runner_)
     return;  // Assume we're in unittests
   base::TimeDelta max_delay =
       base::TimeDelta::FromMinutes(kMaxStartDeviceRetryDelayMinutes);
@@ -193,7 +193,7 @@
   base::TimeDelta retry_delay = later * 2;
   if (retry_delay > max_delay) { retry_delay = max_delay; }
   if (retry_delay < min_delay) { retry_delay = min_delay; }
-  current->PostDelayedTask(
+  task_runner_->PostDelayedTask(
       FROM_HERE,
       base::Bind(&DeviceRegistrationInfo::StartDevice,
                  weak_factory_.GetWeakPtr(), nullptr,
@@ -327,14 +327,12 @@
 
   LOG(INFO) << "Starting notification channel";
 
-  // If no MessageLoop assume we're in unittests.
-  if (!base::MessageLoop::current()) {
-    LOG(INFO) << "No MessageLoop, not starting notification channel";
+  // If no TaskRunner assume we're in unittests.
+  if (!task_runner_) {
+    LOG(INFO) << "No TaskRunner, not starting notification channel";
     return;
   }
 
-  auto task_runner = base::MessageLoop::current()->task_runner();
-
   if (primary_notification_channel_) {
     primary_notification_channel_->Stop();
     primary_notification_channel_.reset();
@@ -348,7 +346,7 @@
   // poll mode.
   const base::TimeDelta pull_interval = config_->polling_period();
   if (!pull_channel_) {
-    pull_channel_.reset(new PullChannel{pull_interval, task_runner});
+    pull_channel_.reset(new PullChannel{pull_interval, task_runner_});
     pull_channel_->Start(this);
   } else {
     pull_channel_->UpdatePullInterval(pull_interval);
@@ -362,7 +360,7 @@
 
   notification_channel_starting_ = true;
   primary_notification_channel_.reset(new XmppChannel{
-      config_->robot_account(), access_token_, task_runner, shill_client_});
+      config_->robot_account(), access_token_, task_runner_, shill_client_});
   primary_notification_channel_->Start(this);
 }
 
@@ -554,7 +552,7 @@
     VLOG(1) << "Cloud request delayed for "
             << cloud_backoff_entry_->GetTimeUntilRelease()
             << " due to backoff policy";
-    base::MessageLoop::current()->PostDelayedTask(
+    task_runner_->PostDelayedTask(
         FROM_HERE,
         base::Bind(&DeviceRegistrationInfo::SendCloudRequest, AsWeakPtr(),
                    data),
@@ -785,7 +783,7 @@
 void DeviceRegistrationInfo::RetryNotifyCommandAborted(
     const std::string& command_id,
     chromeos::ErrorPtr error) {
-  base::MessageLoop::current()->PostDelayedTask(
+  task_runner_->PostDelayedTask(
       FROM_HERE,
       base::Bind(&DeviceRegistrationInfo::NotifyCommandAborted,
                  weak_factory_.GetWeakPtr(),
@@ -902,7 +900,7 @@
     LOG(INFO) << "New command '" << command_instance->GetName()
               << "' arrived, ID: " << command_instance->GetID();
     std::unique_ptr<CommandProxyInterface> cloud_proxy{
-        new CloudCommandProxy(command_instance.get(), this)};
+        new CloudCommandProxy{command_instance.get(), this}};
     command_instance->AddProxy(std::move(cloud_proxy));
     command_manager_->AddCommand(std::move(command_instance));
   }
diff --git a/buffet/device_registration_info.h b/buffet/device_registration_info.h
index bce8c6f..4f12724 100644
--- a/buffet/device_registration_info.h
+++ b/buffet/device_registration_info.h
@@ -13,7 +13,9 @@
 
 #include <base/callback.h>
 #include <base/macros.h>
+#include <base/memory/ref_counted.h>
 #include <base/memory/weak_ptr.h>
+#include <base/single_thread_task_runner.h>
 #include <base/time/time.h>
 #include <base/timer/timer.h>
 #include <chromeos/backoff_entry.h>
@@ -66,6 +68,7 @@
       const std::shared_ptr<StateManager>& state_manager,
       std::unique_ptr<BuffetConfig> config,
       const std::shared_ptr<chromeos::http::Transport>& transport,
+      const scoped_refptr<base::SingleThreadTaskRunner>& task_runner,
       bool notifications_enabled,
       privetd::ShillClient* shill_client);
 
@@ -293,6 +296,7 @@
 
   // HTTP transport used for communications.
   std::shared_ptr<chromeos::http::Transport> transport_;
+  scoped_refptr<base::SingleThreadTaskRunner> task_runner_;
   // Global command manager.
   std::shared_ptr<CommandManager> command_manager_;
   // Device state manager.
diff --git a/buffet/device_registration_info_unittest.cc b/buffet/device_registration_info_unittest.cc
index c8130c3..7a2a281 100644
--- a/buffet/device_registration_info_unittest.cc
+++ b/buffet/device_registration_info_unittest.cc
@@ -176,7 +176,7 @@
     config_ = config.get();
     dev_reg_.reset(new DeviceRegistrationInfo{command_manager_, state_manager_,
                                               std::move(config), transport_,
-                                              true, nullptr});
+                                              nullptr, true, nullptr});
 
     ReloadConfig();
   }
diff --git a/buffet/manager.cc b/buffet/manager.cc
index 88a0b16..ca0f0ab 100644
--- a/buffet/manager.cc
+++ b/buffet/manager.cc
@@ -12,6 +12,7 @@
 #include <base/bind_helpers.h>
 #include <base/json/json_reader.h>
 #include <base/json/json_writer.h>
+#include <base/message_loop/message_loop.h>
 #include <base/time/time.h>
 #include <chromeos/dbus/async_event_sequencer.h>
 #include <chromeos/dbus/exported_object_manager.h>
@@ -89,6 +90,7 @@
   // device info state data unencrypted.
   device_info_.reset(new DeviceRegistrationInfo(
       command_manager_, state_manager_, std::move(config), transport,
+      base::MessageLoop::current()->task_runner(),
       options.xmpp_enabled, shill_client_.get()));
   device_info_->AddOnRegistrationChangedCallback(base::Bind(
       &Manager::OnRegistrationChanged, weak_ptr_factory_.GetWeakPtr()));