arc: Add user notification to request reboot.

Add ArcSnapshotRebootNotification that is shown if inside a snapshot
update interval and user is logged in.
SnapshotRebootController tries to reboot the device once user consents
by pressing a restart button.

BUG=b:170187468
TEST=components_unittests
TEST=unit_tests

Change-Id: I3dcae858e8f889aa4be4fc7d5344359488c9d5fc
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2627523Reviewed-by: Hidehiko Abe <hidehiko@chromium.org>
Reviewed-by: Oleksandr Peletskyi <peletskyi@chromium.org>
Commit-Queue: Polina Bondarenko <pbond@chromium.org>
Auto-Submit: Polina Bondarenko <pbond@chromium.org>
Cr-Commit-Position: refs/heads/master@{#844690}

chrome/app/chromeos_strings.grdp

@@ -4038,6 +4038,15 @@
   <message name="IDS_ARC_POPUP_HELP_LOADING" desc="Loading message of the Arc popup help.">
     Loading...
   </message>
+  <message name="IDS_ARC_SNAPSHOT_RESTART_NOTIFICATION_TITLE" desc="Title of the notification shown to the user when ARC data snapshot must be updated">
+    Restart requested
+  </message>
+  <message name="IDS_ARC_SNAPSHOT_RESTART_NOTIFICATION_MESSAGE" desc="Message in the notification shown to the user when ARC data snapshot must be updated">
+    You Administrator requested to restart your device to update apps. This may take several minutes to complete
+  </message>
+  <message name="IDS_ARC_SNAPSHOT_RESTART_NOTIFICATION_RESTART_BUTTON" desc="Label on button in the notification shown to the user when ARC data snapshot must be updated">
+    Restart
+  </message>
   <message name="IDS_LOW_DISK_NOTIFICATION_TITLE" desc="Title of the notification warning users that they are low on disk space.">
     Device is low on disk space
   </message>

chrome/app/chromeos_strings_grdp/IDS_ARC_SNAPSHOT_RESTART_NOTIFICATION_MESSAGE.png.sha1

+b09211972ca9a49469070e4a9b6f13ec55902db9
\ No newline at end of file

chrome/app/chromeos_strings_grdp/IDS_ARC_SNAPSHOT_RESTART_NOTIFICATION_RESTART_BUTTON.png.sha1

+b09211972ca9a49469070e4a9b6f13ec55902db9
\ No newline at end of file

chrome/app/chromeos_strings_grdp/IDS_ARC_SNAPSHOT_RESTART_NOTIFICATION_TITLE.png.sha1

+b09211972ca9a49469070e4a9b6f13ec55902db9
\ No newline at end of file

chrome/browser/chromeos/BUILD.gn

@@ -635,6 +635,8 @@ source_set("chromeos") {
     "arc/enterprise/arc_enterprise_reporting_service.h",
     "arc/enterprise/arc_force_installed_apps_tracker.cc",
     "arc/enterprise/arc_force_installed_apps_tracker.h",
+    "arc/enterprise/arc_snapshot_reboot_notification_impl.cc",
+    "arc/enterprise/arc_snapshot_reboot_notification_impl.h",
     "arc/enterprise/cert_store/arc_cert_installer.cc",
     "arc/enterprise/cert_store/arc_cert_installer.h",
     "arc/enterprise/cert_store/arc_cert_installer_utils.cc",
@@ -3381,6 +3383,7 @@ source_set("unit_tests") {
     "arc/boot_phase_monitor/arc_boot_phase_monitor_bridge_unittest.cc",
     "arc/enterprise/arc_data_snapshotd_delegate_unittest.cc",
     "arc/enterprise/arc_force_installed_apps_tracker_unittest.cc",
+    "arc/enterprise/arc_snapshot_reboot_notification_impl_unittest.cc",
     "arc/enterprise/cert_store/arc_cert_installer_unittest.cc",
     "arc/extensions/arc_support_message_host_unittest.cc",
     "arc/file_system_watcher/arc_file_system_watcher_service_unittest.cc",

chrome/browser/chromeos/arc/enterprise/arc_data_snapshotd_delegate.cc

@@ -4,6 +4,7 @@
 
 #include "chrome/browser/chromeos/arc/enterprise/arc_data_snapshotd_delegate.h"
 #include "chrome/browser/chromeos/arc/arc_util.h"
+#include "chrome/browser/chromeos/arc/enterprise/arc_snapshot_reboot_notification_impl.h"
 #include "chrome/browser/chromeos/arc/session/arc_session_manager.h"
 
 namespace arc {
@@ -62,6 +63,11 @@ PrefService* ArcDataSnapshotdDelegate::GetProfilePrefService() {
   return arc_session_manager_->profile()->GetPrefs();
 }
 
+std::unique_ptr<ArcSnapshotRebootNotification>
+ArcDataSnapshotdDelegate::CreateRebootNotification() {
+  return std::make_unique<ArcSnapshotRebootNotificationImpl>();
+}
+
 void ArcDataSnapshotdDelegate::OnArcSessionStopped(arc::ArcStopReason reason) {
   switch (reason) {
     case arc::ArcStopReason::SHUTDOWN:

chrome/browser/chromeos/arc/enterprise/arc_data_snapshotd_delegate.h

@@ -8,6 +8,7 @@
 #include "base/callback.h"
 #include "chrome/browser/chromeos/arc/session/arc_session_manager_observer.h"
 #include "components/arc/enterprise/arc_data_snapshotd_manager.h"
+#include "components/arc/enterprise/arc_snapshot_reboot_notification.h"
 #include "components/arc/session/arc_stop_reason.h"
 
 class PrefService;
@@ -31,6 +32,8 @@ class ArcDataSnapshotdDelegate : public ArcDataSnapshotdManager::Delegate,
   void RequestStopArcInstance(
       base::OnceCallback<void(bool)> stopped_callback) override;
   PrefService* GetProfilePrefService() override;
+  std::unique_ptr<ArcSnapshotRebootNotification> CreateRebootNotification()
+      override;
 
   // arc::ArcSessionManagerObserver overrides:
   void OnArcSessionStopped(arc::ArcStopReason reason) override;

chrome/browser/chromeos/arc/enterprise/arc_snapshot_reboot_notification_impl.cc

+// Copyright 2021 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 "chrome/browser/chromeos/arc/enterprise/arc_snapshot_reboot_notification_impl.h"
+
+#include "ash/public/cpp/notification_utils.h"
+#include "chrome/browser/notifications/system_notification_helper.h"
+#include "chrome/grit/generated_resources.h"
+#include "components/vector_icons/vector_icons.h"
+#include "ui/base/l10n/l10n_util.h"
+#include "ui/message_center/public/cpp/notification.h"
+#include "ui/message_center/public/cpp/notification_types.h"
+#include "ui/message_center/public/cpp/notifier_id.h"
+
+namespace arc {
+namespace data_snapshotd {
+
+namespace {
+
+constexpr char kNotificationId[] = "arc_snaposhot_reboot";
+constexpr int kRestartButtonId = 0;
+
+}  // namespace
+
+ArcSnapshotRebootNotificationImpl::ArcSnapshotRebootNotificationImpl() =
+    default;
+
+ArcSnapshotRebootNotificationImpl::~ArcSnapshotRebootNotificationImpl() {
+  Hide();
+}
+
+void ArcSnapshotRebootNotificationImpl::SetUserConsentCallback(
+    const base::RepeatingClosure& callback) {
+  user_consent_callback_ = callback;
+}
+
+void ArcSnapshotRebootNotificationImpl::Show() {
+  if (shown_)
+    return;
+  shown_ = true;
+
+  message_center::NotifierId notifier_id(
+      message_center::NotifierType::SYSTEM_COMPONENT, kNotificationId);
+
+  message_center::RichNotificationData optional_fields;
+  optional_fields.buttons.push_back(
+      message_center::ButtonInfo(l10n_util::GetStringUTF16(
+          IDS_ARC_SNAPSHOT_RESTART_NOTIFICATION_RESTART_BUTTON)));
+
+  auto notification = ash::CreateSystemNotification(
+      message_center::NOTIFICATION_TYPE_SIMPLE, kNotificationId,
+      l10n_util::GetStringUTF16(IDS_ARC_SNAPSHOT_RESTART_NOTIFICATION_TITLE),
+      l10n_util::GetStringUTF16(IDS_ARC_SNAPSHOT_RESTART_NOTIFICATION_MESSAGE),
+      base::string16() /* display_source */, GURL(), notifier_id,
+      optional_fields,
+      new message_center::HandleNotificationClickDelegate(
+          base::Bind(&ArcSnapshotRebootNotificationImpl::HandleClick,
+                     weak_ptr_factory_.GetWeakPtr())),
+      vector_icons::kBusinessIcon,
+      message_center::SystemNotificationWarningLevel::NORMAL);
+  SystemNotificationHelper::GetInstance()->Display(*notification);
+}
+
+void ArcSnapshotRebootNotificationImpl::Hide() {
+  if (!shown_)
+    return;
+  shown_ = false;
+  SystemNotificationHelper::GetInstance()->Close(kNotificationId);
+}
+
+// static
+std::string
+ArcSnapshotRebootNotificationImpl::get_notification_id_for_testing() {
+  return kNotificationId;
+}
+
+// static
+int ArcSnapshotRebootNotificationImpl::get_restart_button_id_for_testing() {
+  return kRestartButtonId;
+}
+
+void ArcSnapshotRebootNotificationImpl::HandleClick(
+    base::Optional<int> button_index) {
+  if (!button_index)
+    return;
+  DCHECK(button_index.value() == kRestartButtonId);
+  Hide();
+
+  DCHECK(!user_consent_callback_.is_null());
+  user_consent_callback_.Run();
+}
+
+}  // namespace data_snapshotd
+}  // namespace arc

chrome/browser/chromeos/arc/enterprise/arc_snapshot_reboot_notification_impl.h

+// Copyright 2021 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 CHROME_BROWSER_CHROMEOS_ARC_ENTERPRISE_ARC_SNAPSHOT_REBOOT_NOTIFICATION_IMPL_H_
+#define CHROME_BROWSER_CHROMEOS_ARC_ENTERPRISE_ARC_SNAPSHOT_REBOOT_NOTIFICATION_IMPL_H_
+
+#include "base/callback.h"
+#include "base/memory/weak_ptr.h"
+#include "components/arc/enterprise/arc_snapshot_reboot_notification.h"
+
+namespace arc {
+namespace data_snapshotd {
+
+// The actual implementation of notification shown to the user when there is a
+// need to reboot a device to update ARC data snapshot.
+class ArcSnapshotRebootNotificationImpl : public ArcSnapshotRebootNotification {
+ public:
+  ArcSnapshotRebootNotificationImpl();
+  ArcSnapshotRebootNotificationImpl(const ArcSnapshotRebootNotificationImpl&) =
+      delete;
+  ~ArcSnapshotRebootNotificationImpl() override;
+
+  ArcSnapshotRebootNotificationImpl& operator=(
+      const ArcSnapshotRebootNotificationImpl&) = delete;
+
+  // ArcSnapshotRebootNotificationImpl overrides:
+  void SetUserConsentCallback(const base::RepeatingClosure& callback) override;
+  void Show() override;
+  void Hide() override;
+
+  static std::string get_notification_id_for_testing();
+  static int get_restart_button_id_for_testing();
+
+ private:
+  void HandleClick(base::Optional<int> button_index);
+
+  base::RepeatingClosure user_consent_callback_;
+  bool shown_ = false;
+
+  base::WeakPtrFactory<ArcSnapshotRebootNotificationImpl> weak_ptr_factory_{
+      this};
+};
+
+}  // namespace data_snapshotd
+}  // namespace arc
+
+#endif  // CHROME_BROWSER_CHROMEOS_ARC_ENTERPRISE_ARC_SNAPSHOT_REBOOT_NOTIFICATION_IMPL_H_

chrome/browser/chromeos/arc/enterprise/arc_snapshot_reboot_notification_impl_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 "chrome/browser/chromeos/arc/enterprise/arc_snapshot_reboot_notification_impl.h"
+#include "base/callback.h"
+#include "base/run_loop.h"
+#include "base/test/bind.h"
+#include "chrome/browser/notifications/notification_display_service_tester.h"
+#include "chrome/browser/notifications/system_notification_helper.h"
+#include "chrome/test/base/testing_browser_process.h"
+#include "content/public/test/browser_task_environment.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace arc {
+namespace data_snapshotd {
+
+class ArcSnapshotRebootNotificationTest : public testing::Test {
+ public:
+  ArcSnapshotRebootNotificationTest() = default;
+  ArcSnapshotRebootNotificationTest(const ArcSnapshotRebootNotificationTest&) =
+      delete;
+  ArcSnapshotRebootNotificationTest& operator=(
+      const ArcSnapshotRebootNotificationTest&) = delete;
+
+  void SetUp() override {
+    TestingBrowserProcess::GetGlobal()->SetSystemNotificationHelper(
+        std::make_unique<SystemNotificationHelper>());
+    tester_ = std::make_unique<NotificationDisplayServiceTester>(nullptr);
+    tester_->SetNotificationAddedClosure(base::BindRepeating(
+        &ArcSnapshotRebootNotificationTest::OnNotificationAdded,
+        base::Unretained(this)));
+    tester_->SetNotificationClosedClosure(base::BindRepeating(
+        &ArcSnapshotRebootNotificationTest::OnNotificationClosed,
+        base::Unretained(this)));
+  }
+
+  void TearDown() override {
+    EXPECT_FALSE(IsNotificationShown());
+    tester_.reset();
+  }
+
+  void ClickOnNotification() {
+    tester_->SimulateClick(
+        NotificationHandler::Type::TRANSIENT,
+        ArcSnapshotRebootNotificationImpl::get_notification_id_for_testing(),
+        base::nullopt, base::nullopt);
+  }
+
+  void ClickOnRestartButton() {
+    tester_->SimulateClick(
+        NotificationHandler::Type::TRANSIENT,
+        ArcSnapshotRebootNotificationImpl::get_notification_id_for_testing(),
+        ArcSnapshotRebootNotificationImpl::get_restart_button_id_for_testing(),
+        base::nullopt);
+  }
+
+  void OnNotificationAdded() { is_notification_shown_ = true; }
+
+  void OnNotificationClosed() { is_notification_shown_ = false; }
+
+  bool IsNotificationShown() { return is_notification_shown_; }
+
+ private:
+  content::BrowserTaskEnvironment task_environment_;
+
+  bool is_notification_shown_ = false;
+
+  std::unique_ptr<NotificationDisplayServiceTester> tester_;
+};
+
+TEST_F(ArcSnapshotRebootNotificationTest, Basic) {
+  ArcSnapshotRebootNotificationImpl notification;
+  EXPECT_FALSE(IsNotificationShown());
+
+  notification.Show();
+  EXPECT_TRUE(IsNotificationShown());
+
+  notification.Hide();
+  EXPECT_FALSE(IsNotificationShown());
+}
+
+TEST_F(ArcSnapshotRebootNotificationTest, ClickOnRestartButton) {
+  ArcSnapshotRebootNotificationImpl notification;
+  notification.Show();
+  EXPECT_TRUE(IsNotificationShown());
+
+  base::RunLoop run_loop;
+  notification.SetUserConsentCallback(run_loop.QuitClosure());
+
+  ClickOnRestartButton();
+  run_loop.Run();
+  EXPECT_FALSE(IsNotificationShown());
+}
+
+TEST_F(ArcSnapshotRebootNotificationTest, ClickOnNotification) {
+  ArcSnapshotRebootNotificationImpl notification;
+  notification.Show();
+  EXPECT_TRUE(IsNotificationShown());
+  notification.SetUserConsentCallback(base::BindLambdaForTesting(
+      []() { NOTREACHED() << "Unexpected user consent registered"; }));
+  ClickOnNotification();
+  EXPECT_TRUE(IsNotificationShown());
+}
+
+TEST_F(ArcSnapshotRebootNotificationTest, DoubleShow) {
+  ArcSnapshotRebootNotificationImpl notification;
+  EXPECT_FALSE(IsNotificationShown());
+
+  notification.Show();
+  notification.Show();
+  EXPECT_TRUE(IsNotificationShown());
+}
+
+TEST_F(ArcSnapshotRebootNotificationTest, DoubleHide) {
+  ArcSnapshotRebootNotificationImpl notification;
+  EXPECT_FALSE(IsNotificationShown());
+
+  notification.Show();
+  EXPECT_TRUE(IsNotificationShown());
+
+  notification.Hide();
+  EXPECT_FALSE(IsNotificationShown());
+  notification.Hide();
+}
+
+}  // namespace data_snapshotd
+}  // namespace arc

components/arc/BUILD.gn

@@ -352,6 +352,8 @@ static_library("arc_test_support") {
     "test/fake_policy_instance.h",
     "test/fake_power_instance.cc",
     "test/fake_power_instance.h",
+    "test/fake_snapshot_reboot_notification.cc",
+    "test/fake_snapshot_reboot_notification.h",
     "test/fake_timer_instance.cc",
     "test/fake_timer_instance.h",
     "test/fake_wake_lock_instance.cc",

components/arc/enterprise/BUILD.gn

@@ -11,6 +11,7 @@ static_library("enterprise") {
     "arc_data_snapshotd_bridge.h",
     "arc_data_snapshotd_manager.cc",
     "arc_data_snapshotd_manager.h",
+    "arc_snapshot_reboot_notification.h",
     "snapshot_hours_policy_service.cc",
     "snapshot_hours_policy_service.h",
     "snapshot_reboot_controller.cc",

components/arc/enterprise/arc_data_snapshotd_manager.cc

@@ -510,7 +510,8 @@ void ArcDataSnapshotdManager::OnSnapshotUpdateEndTimeChanged() {
       snapshot_.Sync();
 
       // Request  device to be reboot in a blocked UI mode.
-      reboot_controller_ = std::make_unique<SnapshotRebootController>();
+      reboot_controller_ = std::make_unique<SnapshotRebootController>(
+          delegate_->CreateRebootNotification());
       return;
     case State::kBlockedUi:
     case State::kMgsToLaunch:

components/arc/enterprise/arc_data_snapshotd_manager.h

@@ -71,6 +71,10 @@ class ArcDataSnapshotdManager final
     // Returns a current profile pref service. Should be called only when ARC
     // session is up and running.
     virtual PrefService* GetProfilePrefService() = 0;
+
+    // Creates a snapshot reboot notification.
+    virtual std::unique_ptr<ArcSnapshotRebootNotification>
+    CreateRebootNotification() = 0;
   };
 
   // This class operates with a snapshot related info either last or

components/arc/enterprise/arc_data_snapshotd_manager_unittest.cc

@@ -20,6 +20,7 @@
 #include "components/arc/enterprise/arc_data_snapshotd_bridge.h"
 #include "components/arc/enterprise/snapshot_session_controller.h"
 #include "components/arc/test/fake_apps_tracker.h"
+#include "components/arc/test/fake_snapshot_reboot_notification.h"
 #include "components/prefs/testing_pref_service.h"
 #include "components/session_manager/core/session_manager.h"
 #include "components/user_manager/fake_user_manager.h"
@@ -72,6 +73,11 @@ class FakeDelegate : public ArcDataSnapshotdManager::Delegate {
 
   PrefService* GetProfilePrefService() override { return &pref_service_; }
 
+  std::unique_ptr<ArcSnapshotRebootNotification> CreateRebootNotification()
+      override {
+    return std::make_unique<FakeSnapshotRebootNotification>();
+  }
+
   bool stopped_callback_num() const { return stopped_callback_num_; }
 
  private:

components/arc/enterprise/arc_snapshot_reboot_notification.h

+// Copyright 2021 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_ARC_SNAPSHOT_REBOOT_NOTIFICATION_H_
+#define COMPONENTS_ARC_ENTERPRISE_ARC_SNAPSHOT_REBOOT_NOTIFICATION_H_
+
+#include "base/callback_forward.h"
+
+namespace arc {
+namespace data_snapshotd {
+
+class ArcSnapshotRebootNotification {
+ public:
+  virtual ~ArcSnapshotRebootNotification() = default;
+
+  // Sets a user consent callback |callback| to be invoked once the user
+  // consents to reboot the device.
+  virtual void SetUserConsentCallback(
+      const base::RepeatingClosure& callback) = 0;
+
+  // Shows a snapshot reboot notification.
+  virtual void Show() = 0;
+
+  // Hides a snapshot reboot notification.
+  virtual void Hide() = 0;
+};
+
+}  // namespace data_snapshotd
+}  // namespace arc
+
+#endif  // COMPONENTS_ARC_ENTERPRISE_ARC_SNAPSHOT_REBOOT_NOTIFICATION_H_

components/arc/enterprise/snapshot_reboot_controller.cc

@@ -31,10 +31,15 @@ bool IsUserLoggedIn() {
 const int kMaxRebootAttempts = 3;
 const base::TimeDelta kRebootAttemptDelay = base::TimeDelta::FromMinutes(5);
 
-SnapshotRebootController::SnapshotRebootController() {
+SnapshotRebootController::SnapshotRebootController(
+    std::unique_ptr<ArcSnapshotRebootNotification> notification)
+    : notification_(std::move(notification)) {
+  notification_->SetUserConsentCallback(
+      base::BindRepeating(&SnapshotRebootController::HandleUserConsent,
+                          weak_ptr_factory_.GetWeakPtr()));
   session_manager::SessionManager::Get()->AddObserver(this);
   if (IsUserLoggedIn()) {
-    // TODO (pbond): show notification.
+    notification_->Show();
   } else {
     // The next operation after reboot is blocking, ensure no one uses device
     // during the next 5 mins.
@@ -49,10 +54,11 @@ SnapshotRebootController::~SnapshotRebootController() {
 void SnapshotRebootController::OnSessionStateChanged() {
   if (IsUserLoggedIn()) {
     StopRebootTimer();
-    // TODO(pbond): show notification.
+    notification_->Show();
   } else {
     // The next operation after reboot is blocking, ensure no one uses device
     // during the next 5 mins.
+    notification_->Hide();
     StartRebootTimer();
   }
 }
@@ -87,5 +93,10 @@ void SnapshotRebootController::OnRebootTimer() {
   SetRebootTimer();
 }
 
+void SnapshotRebootController::HandleUserConsent() {
+  RequestReboot();
+  StartRebootTimer();
+}
+
 }  // namespace data_snapshotd
 }  // namespace arc

components/arc/enterprise/snapshot_reboot_controller.h

@@ -7,6 +7,7 @@
 
 #include "base/memory/weak_ptr.h"
 #include "base/timer/timer.h"
+#include "components/arc/enterprise/arc_snapshot_reboot_notification.h"
 #include "components/session_manager/core/session_manager_observer.h"
 
 namespace arc {
@@ -23,7 +24,8 @@ extern const base::TimeDelta kRebootAttemptDelay;
 class SnapshotRebootController
     : public session_manager::SessionManagerObserver {
  public:
-  SnapshotRebootController();
+  SnapshotRebootController(
+      std::unique_ptr<ArcSnapshotRebootNotification> notification);
   SnapshotRebootController(const SnapshotRebootController&) = delete;
   ~SnapshotRebootController() override;
 
@@ -39,10 +41,15 @@ class SnapshotRebootController
   void SetRebootTimer();
   void StopRebootTimer();
   void OnRebootTimer();
+  // This callback is called once user consents to restart a device.
+  // The device is requested to be restarted immediately.
+  void HandleUserConsent();
 
   base::OneShotTimer reboot_timer_;
   int reboot_attempts_ = 0;
 
+  std::unique_ptr<ArcSnapshotRebootNotification> notification_;
+
   base::WeakPtrFactory<SnapshotRebootController> weak_ptr_factory_{this};
 };
 

components/arc/enterprise/snapshot_reboot_controller_unittest.cc

@@ -7,6 +7,7 @@
 #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/arc/test/fake_snapshot_reboot_notification.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"
@@ -62,7 +63,6 @@ class SnapshotRebootControllerTest : public testing::Test {
   }
 
   user_manager::FakeUserManager* user_manager() { return fake_user_manager_; }
-
  private:
   base::test::TaskEnvironment task_environment_{
       base::test::TaskEnvironment::TimeSource::MOCK_TIME};
@@ -75,31 +75,62 @@ class SnapshotRebootControllerTest : public testing::Test {
 TEST_F(SnapshotRebootControllerTest, UserLoggedIn) {
   LoginUserSession();
   EXPECT_TRUE(user_manager()->IsUserLoggedIn());
-  SnapshotRebootController controller;
+
+  auto notification_ptr = std::make_unique<FakeSnapshotRebootNotification>();
+  auto* notification = notification_ptr.get();
+  SnapshotRebootController controller(std::move(notification_ptr));
+
   EXPECT_FALSE(controller.get_timer_for_testing()->IsRunning());
   EXPECT_EQ(0, client()->num_request_restart_calls());
+  EXPECT_TRUE(notification->shown());
 }
 
 TEST_F(SnapshotRebootControllerTest, BasicReboot) {
-  SnapshotRebootController controller;
+  auto notification_ptr = std::make_unique<FakeSnapshotRebootNotification>();
+  auto* notification = notification_ptr.get();
+  SnapshotRebootController controller(std::move(notification_ptr));
+
+  EXPECT_FALSE(notification->shown());
   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(notification->shown());
   }
   EXPECT_FALSE(controller.get_timer_for_testing()->IsRunning());
 }
 
 TEST_F(SnapshotRebootControllerTest, OnSessionStateChangedLogin) {
-  SnapshotRebootController controller;
+  auto notification_ptr = std::make_unique<FakeSnapshotRebootNotification>();
+  auto* notification = notification_ptr.get();
+  SnapshotRebootController controller(std::move(notification_ptr));
+
   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());
+  EXPECT_TRUE(notification->shown());
 }
 
+TEST_F(SnapshotRebootControllerTest, UserConsentReboot) {
+  LoginUserSession();
+  EXPECT_TRUE(user_manager()->IsUserLoggedIn());
+
+  auto notification_ptr = std::make_unique<FakeSnapshotRebootNotification>();
+  auto* notification = notification_ptr.get();
+  SnapshotRebootController controller(std::move(notification_ptr));
+
+  EXPECT_FALSE(controller.get_timer_for_testing()->IsRunning());
+  EXPECT_EQ(0, client()->num_request_restart_calls());
+  EXPECT_TRUE(notification->shown());
+
+  // The reboot is requested immediately once user consents to reboot.
+  notification->Click();
+  EXPECT_EQ(1, client()->num_request_restart_calls());
+  EXPECT_TRUE(controller.get_timer_for_testing()->IsRunning());
+}
 }  // namespace data_snapshotd
 }  // namespace arc

components/arc/test/fake_snapshot_reboot_notification.cc

+// Copyright 2021 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/test/fake_snapshot_reboot_notification.h"
+
+namespace arc {
+namespace data_snapshotd {
+
+FakeSnapshotRebootNotification::FakeSnapshotRebootNotification() = default;
+FakeSnapshotRebootNotification::~FakeSnapshotRebootNotification() = default;
+
+void FakeSnapshotRebootNotification::SetUserConsentCallback(
+    const base::RepeatingClosure& closure) {
+  user_consent_callback_ = closure;
+}
+
+void FakeSnapshotRebootNotification::Show() {
+  shown_ = true;
+}
+
+void FakeSnapshotRebootNotification::Hide() {
+  shown_ = false;
+}
+
+void FakeSnapshotRebootNotification::Click() {
+  user_consent_callback_.Run();
+}
+
+}  // namespace data_snapshotd
+}  // namespace arc

components/arc/test/fake_snapshot_reboot_notification.h

+// Copyright 2021 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_TEST_FAKE_SNAPSHOT_REBOOT_NOTIFICATION_H_
+#define COMPONENTS_ARC_TEST_FAKE_SNAPSHOT_REBOOT_NOTIFICATION_H_
+
+#include "base/callback.h"
+#include "components/arc/enterprise/arc_snapshot_reboot_notification.h"
+
+namespace arc {
+namespace data_snapshotd {
+
+// Fake implementation of ArcSnapshotRebootNotification for tests.
+class FakeSnapshotRebootNotification : public ArcSnapshotRebootNotification {
+ public:
+  FakeSnapshotRebootNotification();
+  FakeSnapshotRebootNotification(const FakeSnapshotRebootNotification&) =
+      delete;
+  FakeSnapshotRebootNotification& operator=(
+      const FakeSnapshotRebootNotification&) = delete;
+
+  ~FakeSnapshotRebootNotification() override;
+
+  // ArcSnapshotRebootNotification overrides:
+  void SetUserConsentCallback(const base::RepeatingClosure& closure) override;
+  void Show() override;
+  void Hide() override;
+
+  void Click();
+  bool shown() const { return shown_; }
+
+ private:
+  base::RepeatingClosure user_consent_callback_;
+  bool shown_ = false;
+};
+
+}  // namespace data_snapshotd
+}  // namespace arc
+
+#endif  // COMPONENTS_ARC_TEST_FAKE_SNAPSHOT_REBOOT_NOTIFICATION_H_