arc: Request the device's restart to take a snapshot.

Add SnapshotRebootController that is responsible for rebooting the
device if no user is logged in.
SnapshotRebootController makes 3 attempts every 5 mins to reboot a
device if inside a snapshot update interval.

BUG=b:170187468
TEST=components_unittests

Change-Id: Ic19c2386d121e403d42546b25fd67c0c221d52b2
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2584506
Commit-Queue: Polina Bondarenko <pbond@chromium.org>
Auto-Submit: Polina Bondarenko <pbond@chromium.org>
Reviewed-by: Yusuke Sato <yusukes@chromium.org>
Reviewed-by: Oleksandr Peletskyi <peletskyi@chromium.org>
Cr-Commit-Position: refs/heads/master@{#844057}

components/arc/BUILD.gn

@@ -396,6 +396,7 @@ source_set("unit_tests") {
     "enterprise/arc_data_snapshotd_bridge_unittest.cc",
     "enterprise/arc_data_snapshotd_manager_unittest.cc",
     "enterprise/snapshot_hours_policy_service_unittest.cc",
+    "enterprise/snapshot_reboot_controller_unittest.cc",
     "enterprise/snapshot_session_controller_unittest.cc",
     "ime/arc_ime_service_unittest.cc",
     "ime/key_event_result_receiver_unittest.cc",

components/arc/enterprise/BUILD.gn

@@ -13,6 +13,8 @@ static_library("enterprise") {
     "arc_data_snapshotd_manager.h",
     "snapshot_hours_policy_service.cc",
     "snapshot_hours_policy_service.h",
+    "snapshot_reboot_controller.cc",
+    "snapshot_reboot_controller.h",
     "snapshot_session_controller.cc",
     "snapshot_session_controller.h",
   ]
@@ -23,6 +25,7 @@ static_library("enterprise") {
     "//chromeos/cryptohome",
     "//chromeos/dbus:dbus",
     "//chromeos/dbus/arc:arc",
+    "//chromeos/dbus/power",
     "//chromeos/dbus/upstart:upstart",
     "//components/arc:prefs",
     "//components/prefs",

components/arc/enterprise/arc_data_snapshotd_manager.cc

@@ -52,12 +52,33 @@ bool IsRestoredSession() {
          !command_line->HasSwitch(chromeos::switches::kLoginManager);
 }
 
+// Returns true if it is the first Chrome start up after reboot.
+bool IsFirstExecAfterBoot() {
+  return user_manager::UserManager::Get() &&
+         user_manager::UserManager::Get()->IsFirstExecAfterBoot();
+}
+
 // Enables ozone platform headless via command line.
 void EnableHeadlessMode() {
   auto* command_line = base::CommandLine::ForCurrentProcess();
   command_line->AppendSwitchASCII(switches::kOzonePlatform, "headless");
 }
 
+// Returns non-empty account ID string if a MGS is active.
+// Otherwise returns an empty string.
+std::string GetMgsCryptohomeAccountId() {
+  // Take snapshots only for MGSs.
+  if (user_manager::UserManager::Get() &&
+      user_manager::UserManager::Get()->IsLoggedInAsPublicAccount() &&
+      user_manager::UserManager::Get()->GetActiveUser()) {
+    return cryptohome::Identification(user_manager::UserManager::Get()
+                                          ->GetActiveUser()
+                                          ->GetAccountId())
+        .id();
+  }
+  return std::string();
+}
+
 }  // namespace
 
 bool ArcDataSnapshotdManager::is_snapshot_enabled_for_testing_ = false;
@@ -287,7 +308,8 @@ ArcDataSnapshotdManager::ArcDataSnapshotdManager(
   if (IsRestoredSession()) {
     state_ = State::kRestored;
   } else {
-    if (snapshot_.is_blocked_ui_mode() && IsSnapshotEnabled()) {
+    if (snapshot_.is_blocked_ui_mode() && IsSnapshotEnabled() &&
+        IsFirstExecAfterBoot()) {
       state_ = State::kBlockedUi;
       EnableHeadlessMode();
     }
@@ -334,7 +356,7 @@ void ArcDataSnapshotdManager::StartLoadingSnapshot(base::OnceClosure callback) {
     std::move(callback).Run();
     return;
   }
-  std::string account_id = GetCryptohomeAccountId();
+  std::string account_id = GetMgsCryptohomeAccountId();
   if (!account_id.empty() && IsSnapshotEnabled() &&
       (snapshot_.last() || snapshot_.previous())) {
     state_ = State::kLoading;
@@ -439,9 +461,11 @@ void ArcDataSnapshotdManager::OnSnapshotsDisabled() {
     case State::kMgsLaunched:
     case State::kMgsToLaunch:
       state_ = State::kStopping;
+      snapshot_.set_blocked_ui_mode(false);
       if (session_controller_)
         session_controller_->RemoveObserver(this);
       session_controller_.reset();
+      reboot_controller_.reset();
       break;
     // Otherwise, stop all flows, clear snapshots and do not restart browser.
     case State::kNone:
@@ -454,11 +478,48 @@ void ArcDataSnapshotdManager::OnSnapshotsDisabled() {
 }
 
 void ArcDataSnapshotdManager::OnSnapshotUpdateEndTimeChanged() {
-  if (policy_service_.snapshot_update_end_time().is_null())
+  if (policy_service_.snapshot_update_end_time().is_null()) {
+    // Process the end of the snapshot update interval.
+    if (reboot_controller_) {
+      // Stop the reboot process if already requested.
+      snapshot_.set_blocked_ui_mode(false);
+      snapshot_.Sync();
+    }
+    reboot_controller_.reset();
     return;
+  }
   if (!IsSnapshotEnabled())
     return;
-  // TODO(pbond): may be start a reboot process to update a snapshot.
+  // Snapshot can be updated if necessary. Inside the snapshot update interval.
+  // Do not request the reboot of device if already requested.
+  if (reboot_controller_)
+    return;
+  // Do not reboot if last and previous snapshots exist and should not be
+  // updated.
+  if (snapshot_.last() && !snapshot_.last()->updated() &&
+      snapshot_.previous() && !snapshot_.previous()->updated()) {
+    return;
+  }
+
+  switch (state_) {
+    case State::kNone:
+    case State::kLoading:
+    case State::kRestored:
+    case State::kRunning:
+      snapshot_.set_blocked_ui_mode(true);
+      snapshot_.Sync();
+
+      // Request  device to be reboot in a blocked UI mode.
+      reboot_controller_ = std::make_unique<SnapshotRebootController>();
+      return;
+    case State::kBlockedUi:
+    case State::kMgsToLaunch:
+    case State::kMgsLaunched:
+    case State::kStopping:
+      // Do not reboot the device if in blocked UI mode or in process of
+      // disabling the feature.
+      return;
+  }
 }
 
 bool ArcDataSnapshotdManager::IsSnapshotEnabled() {
@@ -663,7 +724,7 @@ void ArcDataSnapshotdManager::OnArcInstanceStopped(bool success) {
     OnSnapshotTaken(false /* success */);
     return;
   }
-  std::string account_id = GetCryptohomeAccountId();
+  std::string account_id = GetMgsCryptohomeAccountId();
   if (account_id.empty() || state_ != State::kMgsLaunched) {
     LOG(ERROR) << "Cryptohome account ID is empty.";
     OnSnapshotTaken(false /* success */);
@@ -738,18 +799,5 @@ void ArcDataSnapshotdManager::OnUiUpdated(bool success) {
     LOG(ERROR) << "Failed to update UI progress bar.";
 }
 
-std::string ArcDataSnapshotdManager::GetCryptohomeAccountId() {
-  // Take snapshots only for MGSs.
-  if (user_manager::UserManager::Get() &&
-      user_manager::UserManager::Get()->IsLoggedInAsPublicAccount() &&
-      user_manager::UserManager::Get()->GetActiveUser()) {
-    return cryptohome::Identification(user_manager::UserManager::Get()
-                                          ->GetActiveUser()
-                                          ->GetAccountId())
-        .id();
-  }
-  return "";
-}
-
 }  // namespace data_snapshotd
 }  // namespace arc

components/arc/enterprise/arc_data_snapshotd_manager.h

@@ -12,6 +12,7 @@
 #include "base/memory/weak_ptr.h"
 #include "components/arc/enterprise/arc_apps_tracker.h"
 #include "components/arc/enterprise/snapshot_hours_policy_service.h"
+#include "components/arc/enterprise/snapshot_reboot_controller.h"
 #include "components/arc/enterprise/snapshot_session_controller.h"
 #include "components/session_manager/core/session_manager_observer.h"
 
@@ -107,6 +108,8 @@ class ArcDataSnapshotdManager final
 
     bool is_last() const { return is_last_; }
 
+    bool updated() const { return updated_; }
+
    private:
     SnapshotInfo(const std::string& os_version,
                  const std::string& creation_date,
@@ -255,6 +258,10 @@ class ArcDataSnapshotdManager final
     session_controller_ = std::move(session_controller);
   }
 
+  SnapshotRebootController* get_reboot_controller_for_testing() const {
+    return reboot_controller_.get();
+  }
+
  private:
   // Attempts to arc-data-snapshotd daemon regardless of state of the class.
   // Runs |callback| once finished.
@@ -307,10 +314,6 @@ class ArcDataSnapshotdManager final
   // Called once a progress bar is updated.
   void OnUiUpdated(bool success);
 
-  // Returns non-empty account ID string if a MGS is active.
-  // Otherwise returns an empty string.
-  std::string GetCryptohomeAccountId();
-
   static bool is_snapshot_enabled_for_testing_;
 
   SnapshotHoursPolicyService policy_service_;
@@ -333,6 +336,9 @@ class ArcDataSnapshotdManager final
   // events.
   std::unique_ptr<SnapshotSessionController> session_controller_;
 
+  // Initialized only when the device reboot is requested.
+  std::unique_ptr<SnapshotRebootController> reboot_controller_;
+
   // Used for cancelling previously posted tasks to daemon.
   base::WeakPtrFactory<ArcDataSnapshotdManager> daemon_weak_ptr_factory_{this};
   // WeakPtrFactory to use for callbacks.

components/arc/enterprise/arc_data_snapshotd_manager_unittest.cc

@@ -147,11 +147,16 @@ class ArcDataSnapshotdManagerBasicTest : public testing::Test {
     upstart_client_ = std::make_unique<TestUpstartClient>();
 
     arc::prefs::RegisterLocalStatePrefs(local_state_.registry());
+
+    base::CommandLine::ForCurrentProcess()->AppendSwitch(
+        chromeos::switches::kFirstExecAfterBoot);
   }
 
   void SetUp() override { SetDBusClientAvailability(true /* is_available */); }
 
   void TearDown() override {
+    ArcDataSnapshotdManager::set_snapshot_enabled_for_testing(
+        false /* enabled */);
     manager_.reset();
     apps_tracker_ = nullptr;
     delegate_ = nullptr;
@@ -689,6 +694,69 @@ TEST_F(ArcDataSnapshotdManagerBasicTest, OnSnapshotSessionFailedTake) {
                  false /* expected_blocked_ui_mode */);
 }
 
+// Test that if the snapshot update interval is not started (end time is null),
+// the device is not rebooted.
+TEST_F(ArcDataSnapshotdManagerBasicTest, OnSnapshotUpdateEndTimeNullFailure) {
+  SetupLocalState(false /* blocked_ui_mode */);
+  ArcDataSnapshotdManager::set_snapshot_enabled_for_testing(true /* enabled */);
+  ExpectStopDaemon(false /* success */);
+  auto* manager = CreateManager(base::DoNothing());
+  manager->OnSnapshotUpdateEndTimeChanged();
+  EXPECT_FALSE(manager->get_reboot_controller_for_testing());
+}
+
+// Test that if snapshot feature is not enabled, the device is not rebooted.
+TEST_F(ArcDataSnapshotdManagerBasicTest,
+       OnSnapshotUpdateEndTimeDisabledFailure) {
+  SetupLocalState(false /* blocked_ui_mode */);
+  auto* manager = CreateManager(base::DoNothing());
+  manager->policy_service()->set_snapshot_update_end_time_for_testing(
+      base::Time::Now());
+  manager->OnSnapshotUpdateEndTimeChanged();
+  EXPECT_FALSE(manager->get_reboot_controller_for_testing());
+}
+
+// Test that if both snapshots exist and no need to update them, the device is
+// not rebooted.
+TEST_F(ArcDataSnapshotdManagerBasicTest, OnSnapshotUpdateEndTimeExistsFailure) {
+  SetupLocalState(false /* blocked_ui_mode */);
+  ArcDataSnapshotdManager::set_snapshot_enabled_for_testing(true /* enabled */);
+  ExpectStopDaemon(false /* success */);
+  auto* manager = CreateManager(base::DoNothing());
+  manager->policy_service()->set_snapshot_update_end_time_for_testing(
+      base::Time::Now());
+  manager->OnSnapshotUpdateEndTimeChanged();
+  EXPECT_FALSE(manager->get_reboot_controller_for_testing());
+}
+
+// Test the end time changed twice in a raw scenario.
+TEST_F(ArcDataSnapshotdManagerBasicTest, OnSnapshotUpdateEndTimeChanged) {
+  ArcDataSnapshotdManager::set_snapshot_enabled_for_testing(true /* enabled */);
+  ExpectStopDaemon(false /* success */);
+  auto* manager = CreateManager(base::DoNothing());
+  manager->policy_service()->set_snapshot_update_end_time_for_testing(
+      base::Time::Now());
+  // Request reboot in blocked UI mode.
+  manager->OnSnapshotUpdateEndTimeChanged();
+  EXPECT_TRUE(manager->get_reboot_controller_for_testing());
+  CheckSnapshots(0 /* expected_snapshots_number */,
+                 true /* expected_blocked_ui_mode */);
+
+  // The reboot is requested above.
+  manager->OnSnapshotUpdateEndTimeChanged();
+  EXPECT_TRUE(manager->get_reboot_controller_for_testing());
+  CheckSnapshots(0 /* expected_snapshots_number */,
+                 true /* expected_blocked_ui_mode */);
+
+  // Stop requesting a reboot if not inside the snapshot update interval.
+  manager->policy_service()->set_snapshot_update_end_time_for_testing(
+      base::Time());
+  manager->OnSnapshotUpdateEndTimeChanged();
+  EXPECT_FALSE(manager->get_reboot_controller_for_testing());
+  CheckSnapshots(0 /* expected_snapshots_number */,
+                 false /* expected_blocked_ui_mode */);
+}
+
 // Test that no one state should lead to any changes except when MGS is expected
 // to be launched.
 TEST_P(ArcDataSnapshotdManagerStateTest, OnSnapshotSessionStarted) {
@@ -774,6 +842,35 @@ TEST_P(ArcDataSnapshotdManagerStateTest, OnSnapshotsDisabled) {
                  false /* expected_blocked_ui_mode */);
 }
 
+TEST_P(ArcDataSnapshotdManagerStateTest, OnSnapshotUpdateEndTimeChanged) {
+  ArcDataSnapshotdManager::set_snapshot_enabled_for_testing(true /* enabled */);
+  ExpectStopDaemon(false /* success */);
+  auto* manager = CreateManager(base::DoNothing());
+  manager->policy_service()->set_snapshot_update_end_time_for_testing(
+      base::Time::Now());
+  manager->set_state_for_testing(expected_state());
+  manager->OnSnapshotUpdateEndTimeChanged();
+
+  switch (expected_state()) {
+    case ArcDataSnapshotdManager::State::kNone:
+    case ArcDataSnapshotdManager::State::kLoading:
+    case ArcDataSnapshotdManager::State::kRestored:
+    case ArcDataSnapshotdManager::State::kRunning:
+      EXPECT_TRUE(manager->get_reboot_controller_for_testing());
+      CheckSnapshots(0 /* expected_snapshots_number */,
+                     true /* expected_blocked_ui_mode */);
+      break;
+    case ArcDataSnapshotdManager::State::kBlockedUi:
+    case ArcDataSnapshotdManager::State::kMgsToLaunch:
+    case ArcDataSnapshotdManager::State::kMgsLaunched:
+    case ArcDataSnapshotdManager::State::kStopping:
+      EXPECT_FALSE(manager->get_reboot_controller_for_testing());
+      CheckSnapshots(0 /* expected_snapshots_number */,
+                     false /* expected_blocked_ui_mode */);
+      break;
+  }
+}
+
 INSTANTIATE_TEST_SUITE_P(
     ArcDataSnapshotdManagerTest,
     ArcDataSnapshotdManagerStateTest,

components/arc/enterprise/snapshot_hours_policy_service.h

@@ -69,6 +69,10 @@ class SnapshotHoursPolicyService {
   }
   const util::WallClockTimer* get_timer_for_testing() const { return &timer_; }
 
+  void set_snapshot_update_end_time_for_testing(base::Time time) {
+    snapshot_update_end_time_ = time;
+  }
+
  private:
   // Processes the policy update: either ArcEnabled and
   // DeviceArcDataSnapshotHours.

components/arc/enterprise/snapshot_reboot_controller.cc

+// Copyright 2020 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 "components/arc/enterprise/snapshot_reboot_controller.h"
+
+#include "base/logging.h"
+#include "chromeos/dbus/power/power_manager_client.h"
+#include "components/session_manager/core/session_manager.h"
+#include "components/user_manager/user_manager.h"
+
+namespace arc {
+namespace data_snapshotd {
+
+namespace {
+
+// Requests reboot. The reboot may not happen immediately.
+void RequestReboot() {
+  chromeos::PowerManagerClient::Get()->RequestRestart(
+      power_manager::REQUEST_RESTART_OTHER, "ARC data snapshot");
+}
+
+// Returns true if any user is logged in.
+bool IsUserLoggedIn() {
+  return user_manager::UserManager::Get() &&
+         user_manager::UserManager::Get()->IsUserLoggedIn();
+}
+
+}  // namespace
+
+const int kMaxRebootAttempts = 3;
+const base::TimeDelta kRebootAttemptDelay = base::TimeDelta::FromMinutes(5);
+
+SnapshotRebootController::SnapshotRebootController() {
+  session_manager::SessionManager::Get()->AddObserver(this);
+  if (IsUserLoggedIn()) {
+    // TODO (pbond): show notification.
+  } else {
+    // The next operation after reboot is blocking, ensure no one uses device
+    // during the next 5 mins.
+    StartRebootTimer();
+  }
+}
+
+SnapshotRebootController::~SnapshotRebootController() {
+  session_manager::SessionManager::Get()->RemoveObserver(this);
+}
+
+void SnapshotRebootController::OnSessionStateChanged() {
+  if (IsUserLoggedIn()) {
+    StopRebootTimer();
+    // TODO(pbond): show notification.
+  } else {
+    // The next operation after reboot is blocking, ensure no one uses device
+    // during the next 5 mins.
+    StartRebootTimer();
+  }
+}
+
+void SnapshotRebootController::StartRebootTimer() {
+  if (reboot_timer_.IsRunning())
+    return;
+  reboot_attempts_ = 0;
+  SetRebootTimer();
+}
+
+void SnapshotRebootController::SetRebootTimer() {
+  reboot_timer_.Start(FROM_HERE, kRebootAttemptDelay,
+                      base::BindOnce(&SnapshotRebootController::OnRebootTimer,
+                                     weak_ptr_factory_.GetWeakPtr()));
+}
+
+void SnapshotRebootController::StopRebootTimer() {
+  reboot_attempts_ = 0;
+  if (!reboot_timer_.IsRunning())
+    return;
+  reboot_timer_.Stop();
+}
+
+void SnapshotRebootController::OnRebootTimer() {
+  reboot_attempts_++;
+  RequestReboot();
+  if (reboot_attempts_ >= kMaxRebootAttempts) {
+    LOG(ERROR) << "The number of reboot attempts exceeded for ARC snapshots.";
+    return;
+  }
+  SetRebootTimer();
+}
+
+}  // namespace data_snapshotd
+}  // namespace arc

components/arc/enterprise/snapshot_reboot_controller.h

+// Copyright 2020 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.
+
+#ifndef COMPONENTS_ARC_ENTERPRISE_SNAPSHOT_REBOOT_CONTROLLER_H_
+#define COMPONENTS_ARC_ENTERPRISE_SNAPSHOT_REBOOT_CONTROLLER_H_
+
+#include "base/memory/weak_ptr.h"
+#include "base/timer/timer.h"
+#include "components/session_manager/core/session_manager_observer.h"
+
+namespace arc {
+namespace data_snapshotd {
+
+// Maximum number of consequent reboot attempts.
+extern const int kMaxRebootAttempts;
+
+// A time delta between reboot attempts.
+extern const base::TimeDelta kRebootAttemptDelay;
+
+// This class observes the MGS state changes and requests a reboot as soon as
+// possible.
+class SnapshotRebootController
+    : public session_manager::SessionManagerObserver {
+ public:
+  SnapshotRebootController();
+  SnapshotRebootController(const SnapshotRebootController&) = delete;
+  ~SnapshotRebootController() override;
+
+  SnapshotRebootController& operator=(const SnapshotRebootController&) = delete;
+
+  // session_manager::SessionManagerObserver overrides:
+  void OnSessionStateChanged() override;
+
+  base::OneShotTimer* get_timer_for_testing() { return &reboot_timer_; }
+
+ private:
+  void StartRebootTimer();
+  void SetRebootTimer();
+  void StopRebootTimer();
+  void OnRebootTimer();
+
+  base::OneShotTimer reboot_timer_;
+  int reboot_attempts_ = 0;
+
+  base::WeakPtrFactory<SnapshotRebootController> weak_ptr_factory_{this};
+};
+
+}  // namespace data_snapshotd
+}  // namespace arc
+
+#endif  // COMPONENTS_ARC_ENTERPRISE_SNAPSHOT_REBOOT_CONTROLLER_H_

components/arc/enterprise/snapshot_reboot_controller_unittest.cc

+// Copyright 2020 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 "components/arc/enterprise/snapshot_reboot_controller.h"
+#include "base/test/task_environment.h"
+#include "chromeos/dbus/dbus_thread_manager.h"
+#include "chromeos/dbus/power/fake_power_manager_client.h"
+#include "components/account_id/account_id.h"
+#include "components/session_manager/core/session_manager.h"
+#include "components/user_manager/fake_user_manager.h"
+#include "components/user_manager/scoped_user_manager.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace arc {
+namespace data_snapshotd {
+
+namespace {
+
+constexpr char kPublicAccountEmail[] = "public@localhost";
+
+chromeos::FakePowerManagerClient* client() {
+  return chromeos::FakePowerManagerClient::Get();
+}
+
+}  // namespace
+
+class SnapshotRebootControllerTest : public testing::Test {
+ public:
+  void SetUp() override {
+    // Initialize fake D-Bus client.
+    chromeos::DBusThreadManager::Initialize();
+    EXPECT_TRUE(chromeos::DBusThreadManager::Get()->IsUsingFakes());
+
+    chromeos::PowerManagerClient::InitializeFake();
+
+    fake_user_manager_ = new user_manager::FakeUserManager();
+    scoped_user_manager_ = std::make_unique<user_manager::ScopedUserManager>(
+        base::WrapUnique(fake_user_manager_));
+  }
+
+  void TearDown() override {
+    chromeos::PowerManagerClient::Shutdown();
+
+    chromeos::DBusThreadManager::Shutdown();
+    scoped_user_manager_.reset();
+    fake_user_manager_ = nullptr;
+  }
+
+  void LoginUserSession() {
+    auto account_id = AccountId::FromUserEmail(kPublicAccountEmail);
+    user_manager()->AddUser(account_id);
+
+    user_manager()->UserLoggedIn(
+        account_id, account_id.GetUserEmail() + "-hash", false, false);
+  }
+
+  void FastForwardAttempt() {
+    task_environment_.FastForwardBy(kRebootAttemptDelay);
+    task_environment_.RunUntilIdle();
+  }
+
+  user_manager::FakeUserManager* user_manager() { return fake_user_manager_; }
+
+ private:
+  base::test::TaskEnvironment task_environment_{
+      base::test::TaskEnvironment::TimeSource::MOCK_TIME};
+
+  session_manager::SessionManager session_manager_;
+  user_manager::FakeUserManager* fake_user_manager_;
+  std::unique_ptr<user_manager::ScopedUserManager> scoped_user_manager_;
+};
+
+TEST_F(SnapshotRebootControllerTest, UserLoggedIn) {
+  LoginUserSession();
+  EXPECT_TRUE(user_manager()->IsUserLoggedIn());
+  SnapshotRebootController controller;
+  EXPECT_FALSE(controller.get_timer_for_testing()->IsRunning());
+  EXPECT_EQ(0, client()->num_request_restart_calls());
+}
+
+TEST_F(SnapshotRebootControllerTest, BasicReboot) {
+  SnapshotRebootController controller;
+  EXPECT_EQ(0, client()->num_request_restart_calls());
+  for (int i = 0; i < kMaxRebootAttempts; i++) {
+    EXPECT_TRUE(controller.get_timer_for_testing()->IsRunning());
+    FastForwardAttempt();
+    EXPECT_EQ(i + 1, client()->num_request_restart_calls());
+  }
+  EXPECT_FALSE(controller.get_timer_for_testing()->IsRunning());
+}
+
+TEST_F(SnapshotRebootControllerTest, OnSessionStateChangedLogin) {
+  SnapshotRebootController controller;
+  EXPECT_TRUE(controller.get_timer_for_testing()->IsRunning());
+
+  LoginUserSession();
+  controller.OnSessionStateChanged();
+  EXPECT_FALSE(controller.get_timer_for_testing()->IsRunning());
+  EXPECT_EQ(0, client()->num_request_restart_calls());
+}
+
+}  // namespace data_snapshotd
+}  // namespace arc