libbrillo: policy: Modify UpdateStaging policy due to change in design

Modify the DeviceUpdateStagingSchedule policy code to account for the
changes to use JSON strings. Add JSONToListValue to remove repeated
logic between new JSON based policies.

BUG=chromium:858621
TEST=cros_workon_make libbrillo --test
CQ-DEPEND=CL:1158225

Change-Id: I01f73980c5309d77bfe04861255d6665cb7e2e11
Reviewed-on: https://chromium-review.googlesource.com/1142480
Commit-Ready: Lutz Justen <ljusten@chromium.org>
Tested-by: Lutz Justen <ljusten@chromium.org>
Tested-by: Adolfo Higueros <adokar@google.com>
Reviewed-by: Dan Erat <derat@chromium.org>

5 files changed, 170 insertions, 155 deletions

diff --git a/policy/device_policy.h b/policy/device_policy.h
index 998cc76..5913d8c 100644
--- a/policy/device_policy.h
+++ b/policy/device_policy.h
@@ -9,6 +9,7 @@
 
 #include <set>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include <base/macros.h>
@@ -52,6 +53,15 @@ class DevicePolicy {
     base::TimeDelta end_time;
   };
 
+  // Identifies a <day, percentage> pair in a staging schedule.
+  struct DayPercentagePair {
+    bool operator==(const DayPercentagePair& other) const {
+      return days == other.days && percentage == other.percentage;
+    }
+    int days;
+    int percentage;
+  };
+
   DevicePolicy();
   virtual ~DevicePolicy();
 
@@ -202,15 +212,17 @@ class DevicePolicy {
   virtual bool GetDisallowedTimeIntervals(
       std::vector<WeeklyTimeInterval>* intervals_out) const = 0;
 
-  // Writes the value of the DeviceUpdateStagingPercentOfFleetPerWeek policy to
+  // Writes the value of the DeviceUpdateStagingSchedule policy to
   // |staging_schedule_out|. Returns true on success.
-  // Values are expected to be mononically increasing in the range of [0, 100].
-  // Each value describes the percentage of the fleet that is expected to
-  // receive an update per week, e.g. [5, 20, 30, 100] means that 5% of devices
-  // should be updated in the first week, 20% should be updated in the second
-  // week, and so on.
+  // The schedule is a list of <days, percentage> pairs. The percentages are
+  // expected to be mononically increasing in the range of [1, 100]. Similarly,
+  // days are expected to be monotonically increasing in the range [1, 28]. Each
+  // pair describes the |percentage| of the fleet that is expected to receive an
+  // update after |days| days after an update was discovered. e.g. [<4, 30>, <8,
+  // 100>] means that 30% of devices should be updated in the first 4 days, and
+  // then 100% should be updated after 8 days.
   virtual bool GetDeviceUpdateStagingSchedule(
-      std::vector<int> *staging_schedule_out) const = 0;
+      std::vector<DayPercentagePair>* staging_schedule_out) const = 0;
 
  private:
   // Verifies that the policy signature is correct.
diff --git a/policy/device_policy_impl.cc b/policy/device_policy_impl.cc
index 0e4b852..76b82a1 100644
--- a/policy/device_policy_impl.cc
+++ b/policy/device_policy_impl.cc
@@ -6,6 +6,8 @@
 
 #include <algorithm>
 #include <memory>
+#include <set>
+#include <string>
 
 #include <base/containers/adapters.h>
 #include <base/files/file_util.h>
@@ -18,9 +20,6 @@
 #include <openssl/evp.h>
 #include <openssl/x509.h>
 
-#include <set>
-#include <string>
-
 #include "bindings/chrome_device_policy.pb.h"
 #include "bindings/device_management_backend.pb.h"
 #include "policy/policy_util.h"
@@ -146,6 +145,28 @@ bool DecodeWeeklyTimeFromValue(const base::DictionaryValue& dict_value,
   return true;
 }
 
+std::unique_ptr<base::ListValue> DecodeListValueFromJSON(
+    const std::string& json_string) {
+  std::string error;
+  std::unique_ptr<base::Value> decoded_json =
+      base::JSONReader::ReadAndReturnError(json_string,
+                                           base::JSON_ALLOW_TRAILING_COMMAS,
+                                           nullptr, &error);
+  if (!decoded_json) {
+    LOG(ERROR) << "Invalid JSON string: " << error;
+    return nullptr;
+  }
+
+  std::unique_ptr<base::ListValue> list_val =
+      base::ListValue::From(std::move(decoded_json));
+  if (!list_val) {
+    LOG(ERROR) << "JSON string is not a list";
+    return nullptr;
+  }
+
+  return list_val;
+}
+
 }  // namespace
 
 DevicePolicyImpl::DevicePolicyImpl()
@@ -503,20 +524,34 @@ bool DevicePolicyImpl::GetUsbDetachableWhitelist(
 }
 
 bool DevicePolicyImpl::GetDeviceUpdateStagingSchedule(
-    std::vector<int> *staging_schedule_out) const {
+    std::vector<DayPercentagePair>* staging_schedule_out) const {
+  staging_schedule_out->clear();
+
   if (!device_policy_.has_auto_update_settings())
     return false;
 
   const em::AutoUpdateSettingsProto &proto =
       device_policy_.auto_update_settings();
 
-  if (proto.staging_percent_of_fleet_per_week_size() == 0)
+  if (!proto.has_staging_schedule())
     return false;
 
-  for (int i = 0; i < proto.staging_percent_of_fleet_per_week_size(); i++) {
-    // Limit the percentage to [0, 100]
-    staging_schedule_out->push_back(
-        std::max(std::min(proto.staging_percent_of_fleet_per_week(i), 100), 0));
+  std::unique_ptr<base::ListValue> list_val =
+      DecodeListValueFromJSON(proto.staging_schedule());
+  if (!list_val)
+    return false;
+
+  for (base::Value* const& pair_value : *list_val) {
+    base::DictionaryValue* day_percentage_pair;
+    if (!pair_value->GetAsDictionary(&day_percentage_pair))
+      return false;
+    int days, percentage;
+    if (!day_percentage_pair->GetInteger("days", &days) ||
+        !day_percentage_pair->GetInteger("percentage", &percentage))
+      return false;
+    // Limit the percentage to [0, 100] and days to [1, 28];
+    staging_schedule_out->push_back({std::max(std::min(days, 28), 1),
+                                     std::max(std::min(percentage, 100), 0)});
   }
 
   return true;
@@ -580,6 +615,8 @@ bool DevicePolicyImpl::GetSecondFactorAuthenticationMode(int* mode_out) const {
 
 bool DevicePolicyImpl::GetDisallowedTimeIntervals(
     std::vector<WeeklyTimeInterval>* intervals_out) const {
+  intervals_out->clear();
+
   if (!device_policy_.has_auto_update_settings()) {
     return false;
   }
@@ -591,27 +628,12 @@ bool DevicePolicyImpl::GetDisallowedTimeIntervals(
     return false;
   }
 
-  // Decode the JSON string
-  std::string error;
-  std::unique_ptr<base::Value> decoded_json =
-      base::JSONReader::ReadAndReturnError(proto.disallowed_time_intervals(),
-                                           base::JSON_ALLOW_TRAILING_COMMAS,
-                                           NULL, &error);
-  if (!decoded_json) {
-    LOG(ERROR) << "Invalid JSON string " << error;
-    return false;
-  }
-
   std::unique_ptr<base::ListValue> list_val =
-      base::ListValue::From(std::move(decoded_json));
-  if (!list_val) {
-    LOG(ERROR) << "JSON string is not a list";
+      DecodeListValueFromJSON(proto.disallowed_time_intervals());
+  if (!list_val)
     return false;
-  }
-
-  intervals_out->clear();
 
-  for (const auto& interval_value : *list_val) {
+  for (base::Value* const& interval_value : *list_val) {
     base::DictionaryValue* interval_dict;
     if (!interval_value->GetAsDictionary(&interval_dict)) {
       LOG(ERROR) << "Invalid JSON string given. Interval is not a dict.";
diff --git a/policy/device_policy_impl.h b/policy/device_policy_impl.h
index 9dbe847..6891312 100644
--- a/policy/device_policy_impl.h
+++ b/policy/device_policy_impl.h
@@ -8,6 +8,7 @@
 #include <memory>
 #include <set>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include <base/files/file_path.h>
@@ -81,7 +82,7 @@ class DevicePolicyImpl : public DevicePolicy {
   bool GetDisallowedTimeIntervals(
       std::vector<WeeklyTimeInterval>* intervals_out) const override;
   bool GetDeviceUpdateStagingSchedule(
-      std::vector<int> *staging_schedule_out) const override;
+      std::vector<DayPercentagePair> *staging_schedule_out) const override;
 
   // Methods that can be used only for testing.
   void set_policy_data_for_testing(
diff --git a/policy/mock_device_policy.h b/policy/mock_device_policy.h
index f26433f..90470e2 100644
--- a/policy/mock_device_policy.h
+++ b/policy/mock_device_policy.h
@@ -7,6 +7,7 @@
 
 #include <set>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include <gmock/gmock.h>
@@ -103,8 +104,8 @@ class MockDevicePolicy : public DevicePolicy {
   MOCK_CONST_METHOD1(GetSecondFactorAuthenticationMode, bool(int*));
   MOCK_CONST_METHOD1(GetDisallowedTimeIntervals,
                      bool(std::vector<WeeklyTimeInterval>*));
-  MOCK_CONST_METHOD1(GetDeviceUpdateStagingSchedule, bool(std::vector<int> *));
-
+  MOCK_CONST_METHOD1(GetDeviceUpdateStagingSchedule,
+                     bool(std::vector<DayPercentagePair>*));
   MOCK_METHOD0(VerifyPolicyFiles, bool(void));
   MOCK_METHOD0(VerifyPolicySignature, bool(void));
 };
diff --git a/policy/tests/device_policy_impl_unittest.cc b/policy/tests/device_policy_impl_unittest.cc
index 939c2b5..37c3916 100644
--- a/policy/tests/device_policy_impl_unittest.cc
+++ b/policy/tests/device_policy_impl_unittest.cc
@@ -16,256 +16,235 @@ using testing::ElementsAre;
 
 namespace policy {
 
+class DevicePolicyImplTest : public testing::Test, public DevicePolicyImpl {
+ protected:
+  void InitializePolicy(const char* device_mode,
+                        const em::ChromeDeviceSettingsProto& proto) {
+    device_policy_.set_policy_for_testing(proto);
+    device_policy_.set_install_attributes_for_testing(
+        std::make_unique<MockInstallAttributesReader>(
+            device_mode, true /* initialized */));
+  }
+
+  DevicePolicyImpl device_policy_;
+};
+
 // Enterprise managed.
-TEST(DevicePolicyImplTest, GetOwner_Managed) {
+TEST_F(DevicePolicyImplTest, GetOwner_Managed) {
   em::PolicyData policy_data;
   policy_data.set_username("user@example.com");
   policy_data.set_management_mode(em::PolicyData::ENTERPRISE_MANAGED);
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_data_for_testing(policy_data);
+  device_policy_.set_policy_data_for_testing(policy_data);
 
   std::string owner("something");
-  EXPECT_TRUE(device_policy.GetOwner(&owner));
+  EXPECT_TRUE(device_policy_.GetOwner(&owner));
   EXPECT_TRUE(owner.empty());
 }
 
 // Consumer owned.
-TEST(DevicePolicyImplTest, GetOwner_Consumer) {
+TEST_F(DevicePolicyImplTest, GetOwner_Consumer) {
   em::PolicyData policy_data;
   policy_data.set_username("user@example.com");
   policy_data.set_management_mode(em::PolicyData::LOCAL_OWNER);
   policy_data.set_request_token("codepath-must-ignore-dmtoken");
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_data_for_testing(policy_data);
+  device_policy_.set_policy_data_for_testing(policy_data);
 
   std::string owner;
-  EXPECT_TRUE(device_policy.GetOwner(&owner));
+  EXPECT_TRUE(device_policy_.GetOwner(&owner));
   EXPECT_EQ("user@example.com", owner);
 }
 
 // Consumer owned, username is missing.
-TEST(DevicePolicyImplTest, GetOwner_ConsumerMissingUsername) {
+TEST_F(DevicePolicyImplTest, GetOwner_ConsumerMissingUsername) {
   em::PolicyData policy_data;
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_data_for_testing(policy_data);
+  device_policy_.set_policy_data_for_testing(policy_data);
 
   std::string owner("something");
-  EXPECT_FALSE(device_policy.GetOwner(&owner));
+  EXPECT_FALSE(device_policy_.GetOwner(&owner));
   EXPECT_EQ("something", owner);
 }
 
 // Enterprise managed, denoted by management_mode.
-TEST(DevicePolicyImplTest, IsEnterpriseManaged_ManagementModeManaged) {
+TEST_F(DevicePolicyImplTest, IsEnterpriseManaged_ManagementModeManaged) {
   em::PolicyData policy_data;
   policy_data.set_management_mode(em::PolicyData::ENTERPRISE_MANAGED);
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_data_for_testing(policy_data);
+  device_policy_.set_policy_data_for_testing(policy_data);
 
-  EXPECT_TRUE(device_policy.IsEnterpriseManaged());
+  EXPECT_TRUE(device_policy_.IsEnterpriseManaged());
 }
 
 // Enterprise managed, fallback to DM token.
-TEST(DevicePolicyImplTest, IsEnterpriseManaged_DMTokenManaged) {
+TEST_F(DevicePolicyImplTest, IsEnterpriseManaged_DMTokenManaged) {
   em::PolicyData policy_data;
   policy_data.set_request_token("abc");
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_data_for_testing(policy_data);
+  device_policy_.set_policy_data_for_testing(policy_data);
 
-  EXPECT_TRUE(device_policy.IsEnterpriseManaged());
+  EXPECT_TRUE(device_policy_.IsEnterpriseManaged());
 }
 
 // Consumer owned, denoted by management_mode.
-TEST(DevicePolicyImplTest, IsEnterpriseManaged_ManagementModeConsumer) {
+TEST_F(DevicePolicyImplTest, IsEnterpriseManaged_ManagementModeConsumer) {
   em::PolicyData policy_data;
   policy_data.set_management_mode(em::PolicyData::LOCAL_OWNER);
   policy_data.set_request_token("codepath-must-ignore-dmtoken");
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_data_for_testing(policy_data);
+  device_policy_.set_policy_data_for_testing(policy_data);
 
-  EXPECT_FALSE(device_policy.IsEnterpriseManaged());
+  EXPECT_FALSE(device_policy_.IsEnterpriseManaged());
 }
 
 // Consumer owned, fallback to interpreting absence of DM token.
-TEST(DevicePolicyImplTest, IsEnterpriseManaged_DMTokenConsumer) {
+TEST_F(DevicePolicyImplTest, IsEnterpriseManaged_DMTokenConsumer) {
   em::PolicyData policy_data;
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_data_for_testing(policy_data);
+  device_policy_.set_policy_data_for_testing(policy_data);
 
-  EXPECT_FALSE(device_policy.IsEnterpriseManaged());
+  EXPECT_FALSE(device_policy_.IsEnterpriseManaged());
 }
 
 // RollbackAllowedMilestones is not set.
-TEST(DevicePolicyImplTest, GetRollbackAllowedMilestones_NotSet) {
-  DevicePolicyImpl device_policy;
-  device_policy.set_install_attributes_for_testing(
+TEST_F(DevicePolicyImplTest, GetRollbackAllowedMilestones_NotSet) {
+  device_policy_.set_install_attributes_for_testing(
       std::make_unique<MockInstallAttributesReader>(
           InstallAttributesReader::kDeviceModeEnterprise, true));
 
   int value = -1;
-  ASSERT_TRUE(device_policy.GetRollbackAllowedMilestones(&value));
+  ASSERT_TRUE(device_policy_.GetRollbackAllowedMilestones(&value));
   EXPECT_EQ(0, value);
 }
 
 // RollbackAllowedMilestones is set to a valid value.
-TEST(DevicePolicyImplTest, GetRollbackAllowedMilestones_Set) {
+TEST_F(DevicePolicyImplTest, GetRollbackAllowedMilestones_Set) {
   em::ChromeDeviceSettingsProto device_policy_proto;
   em::AutoUpdateSettingsProto* auto_update_settings =
       device_policy_proto.mutable_auto_update_settings();
   auto_update_settings->set_rollback_allowed_milestones(3);
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_for_testing(device_policy_proto);
-  device_policy.set_install_attributes_for_testing(
-      std::make_unique<MockInstallAttributesReader>(
-          InstallAttributesReader::kDeviceModeEnterprise, true));
+  InitializePolicy(InstallAttributesReader::kDeviceModeEnterprise,
+                   device_policy_proto);
 
   int value = -1;
-  ASSERT_TRUE(device_policy.GetRollbackAllowedMilestones(&value));
+  ASSERT_TRUE(device_policy_.GetRollbackAllowedMilestones(&value));
   EXPECT_EQ(3, value);
 }
 
 // RollbackAllowedMilestones is set to a valid value, using AD.
-TEST(DevicePolicyImplTest, GetRollbackAllowedMilestones_SetAD) {
+TEST_F(DevicePolicyImplTest, GetRollbackAllowedMilestones_SetAD) {
   em::ChromeDeviceSettingsProto device_policy_proto;
   em::AutoUpdateSettingsProto* auto_update_settings =
       device_policy_proto.mutable_auto_update_settings();
   auto_update_settings->set_rollback_allowed_milestones(3);
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_for_testing(device_policy_proto);
-  device_policy.set_install_attributes_for_testing(
-      std::make_unique<MockInstallAttributesReader>(
-          InstallAttributesReader::kDeviceModeEnterpriseAD, true));
-
+  InitializePolicy(InstallAttributesReader::kDeviceModeEnterpriseAD,
+                   device_policy_proto);
   int value = -1;
-  ASSERT_TRUE(device_policy.GetRollbackAllowedMilestones(&value));
+  ASSERT_TRUE(device_policy_.GetRollbackAllowedMilestones(&value));
   EXPECT_EQ(3, value);
 }
 
 // RollbackAllowedMilestones is set to a valid value, but it's not an enterprise
 // device.
-TEST(DevicePolicyImplTest, GetRollbackAllowedMilestones_SetConsumer) {
+TEST_F(DevicePolicyImplTest, GetRollbackAllowedMilestones_SetConsumer) {
   em::ChromeDeviceSettingsProto device_policy_proto;
   em::AutoUpdateSettingsProto* auto_update_settings =
       device_policy_proto.mutable_auto_update_settings();
   auto_update_settings->set_rollback_allowed_milestones(3);
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_for_testing(device_policy_proto);
-  device_policy.set_install_attributes_for_testing(
-      std::make_unique<MockInstallAttributesReader>(
-          InstallAttributesReader::kDeviceModeConsumer, true));
+  InitializePolicy(InstallAttributesReader::kDeviceModeConsumer,
+                   device_policy_proto);
 
   int value = -1;
-  ASSERT_FALSE(device_policy.GetRollbackAllowedMilestones(&value));
+  ASSERT_FALSE(device_policy_.GetRollbackAllowedMilestones(&value));
 }
 
 // RollbackAllowedMilestones is set to an invalid value.
-TEST(DevicePolicyImplTest, GetRollbackAllowedMilestones_SetTooLarge) {
+TEST_F(DevicePolicyImplTest, GetRollbackAllowedMilestones_SetTooLarge) {
   em::ChromeDeviceSettingsProto device_policy_proto;
   em::AutoUpdateSettingsProto* auto_update_settings =
       device_policy_proto.mutable_auto_update_settings();
   auto_update_settings->set_rollback_allowed_milestones(10);
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_for_testing(device_policy_proto);
-  device_policy.set_install_attributes_for_testing(
-      std::make_unique<MockInstallAttributesReader>(
-          InstallAttributesReader::kDeviceModeEnterprise, true));
+  InitializePolicy(InstallAttributesReader::kDeviceModeEnterprise,
+                   device_policy_proto);
 
   int value = -1;
-  ASSERT_TRUE(device_policy.GetRollbackAllowedMilestones(&value));
+  ASSERT_TRUE(device_policy_.GetRollbackAllowedMilestones(&value));
   EXPECT_EQ(4, value);
 }
 
 // RollbackAllowedMilestones is set to an invalid value.
-TEST(DevicePolicyImplTest, GetRollbackAllowedMilestones_SetTooSmall) {
+TEST_F(DevicePolicyImplTest, GetRollbackAllowedMilestones_SetTooSmall) {
   em::ChromeDeviceSettingsProto device_policy_proto;
   em::AutoUpdateSettingsProto* auto_update_settings =
       device_policy_proto.mutable_auto_update_settings();
   auto_update_settings->set_rollback_allowed_milestones(-1);
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_for_testing(device_policy_proto);
-  device_policy.set_install_attributes_for_testing(
-      std::make_unique<MockInstallAttributesReader>(
-          InstallAttributesReader::kDeviceModeEnterprise, true));
+  InitializePolicy(InstallAttributesReader::kDeviceModeEnterprise,
+                   device_policy_proto);
 
   int value = -1;
-  ASSERT_TRUE(device_policy.GetRollbackAllowedMilestones(&value));
+  ASSERT_TRUE(device_policy_.GetRollbackAllowedMilestones(&value));
   EXPECT_EQ(0, value);
 }
 
 // Update staging schedule has no values
-TEST(DevicePolicyImplTest, GetDeviceUpdateStagingSchedule_NoValues) {
+TEST_F(DevicePolicyImplTest, GetDeviceUpdateStagingSchedule_NoValues) {
   em::ChromeDeviceSettingsProto device_policy_proto;
   em::AutoUpdateSettingsProto *auto_update_settings =
       device_policy_proto.mutable_auto_update_settings();
-  auto_update_settings->clear_staging_percent_of_fleet_per_week();
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_for_testing(device_policy_proto);
-  device_policy.set_install_attributes_for_testing(
-      std::make_unique<MockInstallAttributesReader>(
-          InstallAttributesReader::kDeviceModeEnterprise, true));
+  auto_update_settings->set_staging_schedule("[]");
+  InitializePolicy(InstallAttributesReader::kDeviceModeEnterprise,
+                   device_policy_proto);
 
-  std::vector<int> staging_schedule;
-  EXPECT_FALSE(device_policy.GetDeviceUpdateStagingSchedule(&staging_schedule));
+  std::vector<DayPercentagePair> staging_schedule;
+  ASSERT_TRUE(device_policy_.GetDeviceUpdateStagingSchedule(&staging_schedule));
+  EXPECT_EQ(0, staging_schedule.size());
 }
 
 // Update staging schedule has valid values
-TEST(DevicePolicyImplTest, GetDeviceUpdateStagingSchedule_Valid) {
+TEST_F(DevicePolicyImplTest, GetDeviceUpdateStagingSchedule_Valid) {
   em::ChromeDeviceSettingsProto device_policy_proto;
   em::AutoUpdateSettingsProto *auto_update_settings =
       device_policy_proto.mutable_auto_update_settings();
-  auto_update_settings->add_staging_percent_of_fleet_per_week(10);
-  auto_update_settings->add_staging_percent_of_fleet_per_week(20);
-  auto_update_settings->add_staging_percent_of_fleet_per_week(40);
-  auto_update_settings->add_staging_percent_of_fleet_per_week(100);
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_for_testing(device_policy_proto);
-  device_policy.set_install_attributes_for_testing(
-      std::make_unique<MockInstallAttributesReader>(
-          InstallAttributesReader::kDeviceModeEnterprise, true));
+  auto_update_settings->set_staging_schedule(
+      "[{\"days\": 4, \"percentage\": 40}, {\"days\": 10, \"percentage\": "
+      "100}]");
+  InitializePolicy(InstallAttributesReader::kDeviceModeEnterprise,
+                   device_policy_proto);
 
-  std::vector<int> staging_schedule;
-  ASSERT_TRUE(device_policy.GetDeviceUpdateStagingSchedule(&staging_schedule));
-  EXPECT_THAT(staging_schedule, ElementsAre(10, 20, 40, 100));
+  std::vector<DayPercentagePair> staging_schedule;
+  ASSERT_TRUE(device_policy_.GetDeviceUpdateStagingSchedule(&staging_schedule));
+  EXPECT_THAT(staging_schedule, ElementsAre(DayPercentagePair{4, 40},
+                                            DayPercentagePair{10, 100}));
 }
 
 // Update staging schedule has valid values, set using AD.
-TEST(DevicePolicyImplTest, GetDeviceUpdateStagingSchedule_Valid_AD) {
+TEST_F(DevicePolicyImplTest, GetDeviceUpdateStagingSchedule_Valid_AD) {
   em::ChromeDeviceSettingsProto device_policy_proto;
   em::AutoUpdateSettingsProto *auto_update_settings =
       device_policy_proto.mutable_auto_update_settings();
-  auto_update_settings->add_staging_percent_of_fleet_per_week(10);
-  auto_update_settings->add_staging_percent_of_fleet_per_week(20);
-  auto_update_settings->add_staging_percent_of_fleet_per_week(40);
-  auto_update_settings->add_staging_percent_of_fleet_per_week(100);
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_for_testing(device_policy_proto);
-  device_policy.set_install_attributes_for_testing(
-      std::make_unique<MockInstallAttributesReader>(
-          InstallAttributesReader::kDeviceModeEnterpriseAD, true));
+  auto_update_settings->set_staging_schedule(
+      "[{\"days\": 4, \"percentage\": 40}, {\"days\": 10, \"percentage\": "
+      "100}]");
+  InitializePolicy(InstallAttributesReader::kDeviceModeEnterpriseAD,
+                   device_policy_proto);
 
-  std::vector<int> staging_schedule;
-  ASSERT_TRUE(device_policy.GetDeviceUpdateStagingSchedule(&staging_schedule));
-  EXPECT_THAT(staging_schedule, ElementsAre(10, 20, 40, 100));
+  std::vector<DayPercentagePair> staging_schedule;
+  ASSERT_TRUE(device_policy_.GetDeviceUpdateStagingSchedule(&staging_schedule));
+  EXPECT_THAT(staging_schedule, ElementsAre(DayPercentagePair{4, 40},
+                                            DayPercentagePair{10, 100}));
 }
 
 // Update staging schedule has values with values set larger than the max
-// allowed percentage and smaller than the min allowed percentage.
-TEST(DevicePolicyImplTest, GetDeviceUpdateStagingSchedule_SetOutsideAllowable) {
+// allowed days/percentage and smaller than the min allowed days/percentage.
+TEST_F(DevicePolicyImplTest,
+       GetDeviceUpdateStagingSchedule_SetOutsideAllowable) {
   em::ChromeDeviceSettingsProto device_policy_proto;
   em::AutoUpdateSettingsProto *auto_update_settings =
       device_policy_proto.mutable_auto_update_settings();
-  auto_update_settings->add_staging_percent_of_fleet_per_week(-10);
-  auto_update_settings->add_staging_percent_of_fleet_per_week(20);
-  auto_update_settings->add_staging_percent_of_fleet_per_week(40);
-  auto_update_settings->add_staging_percent_of_fleet_per_week(110);
-  DevicePolicyImpl device_policy;
-  device_policy.set_policy_for_testing(device_policy_proto);
-  device_policy.set_install_attributes_for_testing(
-      std::make_unique<MockInstallAttributesReader>(
-          InstallAttributesReader::kDeviceModeEnterprise, true));
-
-  std::vector<int> staging_schedule;
-  ASSERT_TRUE(device_policy.GetDeviceUpdateStagingSchedule(&staging_schedule));
-  EXPECT_THAT(staging_schedule, ElementsAre(0, 20, 40, 100));
+  auto_update_settings->set_staging_schedule(
+      "[{\"days\": -1, \"percentage\": -10}, {\"days\": 30, \"percentage\": "
+      "110}]");
+  InitializePolicy(InstallAttributesReader::kDeviceModeEnterprise,
+                   device_policy_proto);
+
+  std::vector<DayPercentagePair> staging_schedule;
+  ASSERT_TRUE(device_policy_.GetDeviceUpdateStagingSchedule(&staging_schedule));
+  EXPECT_THAT(staging_schedule, ElementsAre(DayPercentagePair{1, 0},
+                                            DayPercentagePair{28, 100}));
 }
 
 }  // namespace policy