libbrillo: Set default values for RollbackAllowedMilestones policy.

RollbackAllowedMilestones policy values must be between 0 and 4.
Default for enterprise devices is 4, for consumer devices the
value is always 0.

Because DevicePolicy is not available for consumer devices, this
CL adds IsConsumerDevice to PolicyProvider, so update_engine can
check whether it should use the consumer 0 or not set anything if
the policy is not yet loaded.

Not setting anything in the case when the device is still in
OOBE is important because setting a smaller value can permanently
disable rollback to an older version.

BUG=chromium:813036
TEST='cros_run_unit_tests --board=cyan --packages libbrillo'

Change-Id: I228b12ddf3ba05f3230d111ac05d74c0330999c4
Reviewed-on: https://chromium-review.googlesource.com/1013461
Commit-Ready: Marton Hunyady <hunyadym@chromium.org>
Tested-by: Marton Hunyady <hunyadym@chromium.org>
Reviewed-by: Dan Erat <derat@chromium.org>

9 files changed, 304 insertions, 102 deletions

diff --git a/install_attributes/mock_install_attributes_reader.cc b/install_attributes/mock_install_attributes_reader.cc
index 6046085..da18988 100644
--- a/install_attributes/mock_install_attributes_reader.cc
+++ b/install_attributes/mock_install_attributes_reader.cc
@@ -14,3 +14,9 @@ MockInstallAttributesReader::MockInstallAttributesReader(
   }
   initialized_ = true;
 }
+
+MockInstallAttributesReader::MockInstallAttributesReader(
+    const std::string& device_mode, bool initialized) {
+  attributes_[kAttrMode] = device_mode;
+  initialized_ = initialized;
+}
diff --git a/install_attributes/mock_install_attributes_reader.h b/install_attributes/mock_install_attributes_reader.h
index 1545ea6..5ccee02 100644
--- a/install_attributes/mock_install_attributes_reader.h
+++ b/install_attributes/mock_install_attributes_reader.h
@@ -13,6 +13,7 @@ class MockInstallAttributesReader : public InstallAttributesReader {
  public:
   explicit MockInstallAttributesReader(
       const cryptohome::SerializedInstallAttributes& install_attributes);
+  MockInstallAttributesReader(const std::string& device_mode, bool initialized);
 };
 
 #endif  // LIBBRILLO_INSTALL_ATTRIBUTES_MOCK_INSTALL_ATTRIBUTES_READER_H_
diff --git a/policy/device_policy_impl.cc b/policy/device_policy_impl.cc
index e6ed0e8..548a372 100644
--- a/policy/device_policy_impl.cc
+++ b/policy/device_policy_impl.cc
@@ -26,6 +26,9 @@ namespace em = enterprise_management;
 
 namespace policy {
 
+// Maximum value of RollbackAllowedMilestones policy.
+const int kMaxRollbackAllowedMilestones = 4;
+
 namespace {
 const char kPolicyPath[] = "/var/lib/whitelist/policy";
 const char kPublicKeyPath[] = "/var/lib/whitelist/owner.key";
@@ -301,15 +304,34 @@ bool DevicePolicyImpl::GetRollbackToTargetVersion(
 
 bool DevicePolicyImpl::GetRollbackAllowedMilestones(
     int* rollback_allowed_milestones) const {
-  if (!device_policy_.has_auto_update_settings())
-    return false;
-
-  const em::AutoUpdateSettingsProto& proto =
-      device_policy_.auto_update_settings();
-  if (!proto.has_rollback_allowed_milestones())
-    return false;
-
-  *rollback_allowed_milestones = proto.rollback_allowed_milestones();
+  // This policy can be only set for devices which are enterprise enrolled.
+  if (!install_attributes_reader_->IsLocked())
+    return false;
+  if (install_attributes_reader_->GetAttribute(
+          InstallAttributesReader::kAttrMode) !=
+          InstallAttributesReader::kDeviceModeEnterprise &&
+      install_attributes_reader_->GetAttribute(
+          InstallAttributesReader::kAttrMode) !=
+          InstallAttributesReader::kDeviceModeEnterpriseAD)
+    return false;
+
+  if (device_policy_.has_auto_update_settings()) {
+    const em::AutoUpdateSettingsProto& proto =
+        device_policy_.auto_update_settings();
+    if (proto.has_rollback_allowed_milestones()) {
+      // Policy is set, enforce minimum and maximum constraints.
+      *rollback_allowed_milestones = proto.rollback_allowed_milestones();
+      if (*rollback_allowed_milestones < 0)
+        *rollback_allowed_milestones = 0;
+      if (*rollback_allowed_milestones > kMaxRollbackAllowedMilestones)
+        *rollback_allowed_milestones = kMaxRollbackAllowedMilestones;
+      return true;
+    }
+  }
+  // Policy is not present, use default for enterprise devices.
+  VLOG(1) << "RollbackAllowedMilestones policy is not set, using default "
+          << kMaxRollbackAllowedMilestones << ".";
+  *rollback_allowed_milestones = kMaxRollbackAllowedMilestones;
   return true;
 }
 
diff --git a/policy/device_policy_impl.h b/policy/device_policy_impl.h
index 2a8d7bb..9ec39ac 100644
--- a/policy/device_policy_impl.h
+++ b/policy/device_policy_impl.h
@@ -91,6 +91,10 @@ class DevicePolicyImpl : public DevicePolicy {
       std::unique_ptr<InstallAttributesReader> install_attributes_reader) {
     install_attributes_reader_ = std::move(install_attributes_reader);
   }
+  void set_policy_for_testing(
+      const enterprise_management::ChromeDeviceSettingsProto& device_policy) {
+    device_policy_ = device_policy;
+  }
   void set_policy_path_for_testing(const base::FilePath& policy_path) {
     policy_path_ = policy_path;
   }
diff --git a/policy/libpolicy.cc b/policy/libpolicy.cc
index 9864cdf..a4fc380 100644
--- a/policy/libpolicy.cc
+++ b/policy/libpolicy.cc
@@ -18,12 +18,14 @@ namespace policy {
 PolicyProvider::PolicyProvider() {
 #ifndef __ANDROID__
   device_policy_ = std::make_unique<DevicePolicyImpl>();
+  install_attributes_reader_ = std::make_unique<InstallAttributesReader>();
 #endif
 }
 
 PolicyProvider::PolicyProvider(std::unique_ptr<DevicePolicy> device_policy)
     : device_policy_(std::move(device_policy)),
-      device_policy_is_loaded_(true) {}
+      device_policy_is_loaded_(true),
+      install_attributes_reader_(std::make_unique<InstallAttributesReader>()) {}
 
 PolicyProvider::~PolicyProvider() {}
 
@@ -47,4 +49,25 @@ const DevicePolicy& PolicyProvider::GetDevicePolicy() const {
   return *device_policy_;
 }
 
+bool PolicyProvider::IsConsumerDevice() const {
+  if (!install_attributes_reader_->IsLocked())
+    return false;
+
+  const std::string& device_mode = install_attributes_reader_->GetAttribute(
+      InstallAttributesReader::kAttrMode);
+  return device_mode != InstallAttributesReader::kDeviceModeEnterprise &&
+      device_mode != InstallAttributesReader::kDeviceModeEnterpriseAD;
+}
+
+void PolicyProvider::SetDevicePolicyForTesting(
+    std::unique_ptr<DevicePolicy> device_policy) {
+  device_policy_ = std::move(device_policy);
+  device_policy_is_loaded_ = true;
+}
+
+void PolicyProvider::SetInstallAttributesReaderForTesting(
+    std::unique_ptr<InstallAttributesReader> install_attributes_reader) {
+  install_attributes_reader_ = std::move(install_attributes_reader);
+}
+
 }  // namespace policy
diff --git a/policy/libpolicy.h b/policy/libpolicy.h
index ac39bd9..5e87f58 100644
--- a/policy/libpolicy.h
+++ b/policy/libpolicy.h
@@ -10,6 +10,8 @@
 
 #include <base/macros.h>
 
+#include "install_attributes/libinstallattributes.h"
+
 #pragma GCC visibility push(default)
 
 namespace policy {
@@ -38,9 +40,20 @@ class PolicyProvider {
   // Returns a value from the device policy cache.
   virtual const DevicePolicy& GetDevicePolicy() const;
 
+  // Returns true if the device is not an enterprise enrolled device, so it
+  // won't have device policy before the next powerwash. Returns false if device
+  // is still in OOBE (so device mode is not determined yet).
+  virtual bool IsConsumerDevice() const;
+
+  void SetDevicePolicyForTesting(
+      std::unique_ptr<DevicePolicy> device_policy);
+  void SetInstallAttributesReaderForTesting(
+      std::unique_ptr<InstallAttributesReader> install_attributes_reader);
+
  private:
   std::unique_ptr<DevicePolicy> device_policy_;
   bool device_policy_is_loaded_ = false;
+  std::unique_ptr<InstallAttributesReader> install_attributes_reader_;
 
   DISALLOW_COPY_AND_ASSIGN(PolicyProvider);
 };
diff --git a/policy/mock_libpolicy.h b/policy/mock_libpolicy.h
index acb9b8b..a0f6920 100644
--- a/policy/mock_libpolicy.h
+++ b/policy/mock_libpolicy.h
@@ -23,6 +23,7 @@ class MockPolicyProvider : public PolicyProvider {
   MOCK_METHOD0(Reload, bool(void));
   MOCK_CONST_METHOD0(device_policy_is_loaded, bool(void));
   MOCK_CONST_METHOD0(GetDevicePolicy, const DevicePolicy&(void));
+  MOCK_CONST_METHOD0(IsConsumerDevice, bool(void));
 
  private:
   DISALLOW_COPY_AND_ASSIGN(MockPolicyProvider);
diff --git a/policy/tests/device_policy_impl_unittest.cc b/policy/tests/device_policy_impl_unittest.cc
index e443394..bc6aae0 100644
--- a/policy/tests/device_policy_impl_unittest.cc
+++ b/policy/tests/device_policy_impl_unittest.cc
@@ -7,6 +7,7 @@
 #include "policy/device_policy_impl.h"
 
 #include "bindings/chrome_device_policy.pb.h"
+#include "install_attributes/mock_install_attributes_reader.h"
 
 namespace em = enterprise_management;
 
@@ -90,4 +91,101 @@ TEST(DevicePolicyImplTest, IsEnterpriseManaged_DMTokenConsumer) {
   EXPECT_FALSE(device_policy.IsEnterpriseManaged());
 }
 
+// RollbackAllowedMilestones is not set.
+TEST(DevicePolicyImplTest, GetRollbackAllowedMilestones_NotSet) {
+  DevicePolicyImpl device_policy;
+  device_policy.set_install_attributes_for_testing(
+      std::make_unique<MockInstallAttributesReader>(
+          InstallAttributesReader::kDeviceModeEnterprise, true));
+
+  int value = -1;
+  ASSERT_TRUE(device_policy.GetRollbackAllowedMilestones(&value));
+  EXPECT_EQ(4, value);
+}
+
+// RollbackAllowedMilestones is set to a valid value.
+TEST(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));
+
+  int value = -1;
+  ASSERT_TRUE(device_policy.GetRollbackAllowedMilestones(&value));
+  EXPECT_EQ(3, value);
+}
+
+// RollbackAllowedMilestones is set to a valid value, using AD.
+TEST(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));
+
+  int value = -1;
+  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) {
+  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));
+
+  int value = -1;
+  ASSERT_FALSE(device_policy.GetRollbackAllowedMilestones(&value));
+}
+
+// RollbackAllowedMilestones is set to an invalid value.
+TEST(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));
+
+  int value = -1;
+  ASSERT_TRUE(device_policy.GetRollbackAllowedMilestones(&value));
+  EXPECT_EQ(4, value);
+}
+
+// RollbackAllowedMilestones is set to an invalid value.
+TEST(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));
+
+  int value = -1;
+  ASSERT_TRUE(device_policy.GetRollbackAllowedMilestones(&value));
+  EXPECT_EQ(0, value);
+}
+
 }  // namespace policy
diff --git a/policy/tests/libpolicy_unittest.cc b/policy/tests/libpolicy_unittest.cc
index 65caaad..310ef00 100644
--- a/policy/tests/libpolicy_unittest.cc
+++ b/policy/tests/libpolicy_unittest.cc
@@ -46,10 +46,11 @@ std::unique_ptr<DevicePolicyImpl> CreateDevicePolicyImpl(
 TEST(PolicyTest, DevicePolicyAllSetTest) {
   base::FilePath policy_file(kPolicyFileAllSet);
   base::FilePath key_file(kKeyFile);
-  PolicyProvider provider(
-      CreateDevicePolicyImpl(std::make_unique<MockInstallAttributesReader>(
-                                 cryptohome::SerializedInstallAttributes()),
-                             policy_file, key_file, false));
+  PolicyProvider provider;
+  provider.SetDevicePolicyForTesting(CreateDevicePolicyImpl(
+      std::make_unique<MockInstallAttributesReader>(
+          InstallAttributesReader::kDeviceModeEnterprise, true),
+      policy_file, key_file, false));
   provider.Reload();
 
   // Ensure we successfully loaded the device policy file.
@@ -60,123 +61,125 @@ TEST(PolicyTest, DevicePolicyAllSetTest) {
   // Check that we can read out all fields of the sample protobuf.
   int int_value = -1;
   ASSERT_TRUE(policy.GetPolicyRefreshRate(&int_value));
-  ASSERT_EQ(100, int_value);
+  EXPECT_EQ(100, int_value);
 
   std::vector<std::string> list_value;
   ASSERT_TRUE(policy.GetUserWhitelist(&list_value));
   ASSERT_EQ(3, list_value.size());
-  ASSERT_EQ("me@here.com", list_value[0]);
-  ASSERT_EQ("you@there.com", list_value[1]);
-  ASSERT_EQ("*@monsters.com", list_value[2]);
+  EXPECT_EQ("me@here.com", list_value[0]);
+  EXPECT_EQ("you@there.com", list_value[1]);
+  EXPECT_EQ("*@monsters.com", list_value[2]);
 
   bool bool_value = true;
   ASSERT_TRUE(policy.GetGuestModeEnabled(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetCameraEnabled(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetShowUserNames(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetDataRoamingEnabled(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetAllowNewUsers(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetMetricsEnabled(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetReportVersionInfo(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetReportActivityTimes(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetReportBootMode(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetEphemeralUsersEnabled(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   std::string string_value;
   ASSERT_TRUE(policy.GetReleaseChannel(&string_value));
-  ASSERT_EQ("stable-channel", string_value);
+  EXPECT_EQ("stable-channel", string_value);
 
   bool_value = false;
   ASSERT_TRUE(policy.GetReleaseChannelDelegated(&bool_value));
-  ASSERT_TRUE(bool_value);
+  EXPECT_TRUE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetUpdateDisabled(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   int64_t int64_value = -1LL;
   ASSERT_TRUE(policy.GetScatterFactorInSeconds(&int64_value));
-  ASSERT_EQ(17LL, int64_value);
+  EXPECT_EQ(17LL, int64_value);
 
   ASSERT_TRUE(policy.GetTargetVersionPrefix(&string_value));
-  ASSERT_EQ("42.0.", string_value);
+  EXPECT_EQ("42.0.", string_value);
+
   int_value = -1;
   ASSERT_TRUE(policy.GetRollbackToTargetVersion(&int_value));
-  ASSERT_EQ(enterprise_management::AutoUpdateSettingsProto::
+  EXPECT_EQ(enterprise_management::AutoUpdateSettingsProto::
                 ROLLBACK_WITH_FULL_POWERWASH,
             int_value);
+
   int_value = -1;
   ASSERT_TRUE(policy.GetRollbackAllowedMilestones(&int_value));
-  ASSERT_EQ(3, int_value);
+  EXPECT_EQ(3, int_value);
 
   std::set<std::string> types;
   ASSERT_TRUE(policy.GetAllowedConnectionTypesForUpdate(&types));
-  ASSERT_TRUE(types.end() != types.find("ethernet"));
-  ASSERT_TRUE(types.end() != types.find("wifi"));
-  ASSERT_EQ(2, types.size());
+  EXPECT_TRUE(types.end() != types.find("ethernet"));
+  EXPECT_TRUE(types.end() != types.find("wifi"));
+  EXPECT_EQ(2, types.size());
 
   ASSERT_TRUE(policy.GetOpenNetworkConfiguration(&string_value));
-  ASSERT_EQ("{}", string_value);
+  EXPECT_EQ("{}", string_value);
 
   ASSERT_TRUE(policy.GetOwner(&string_value));
-  ASSERT_EQ("", string_value);
+  EXPECT_EQ("", string_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetHttpDownloadsEnabled(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetAuP2PEnabled(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   bool_value = true;
   ASSERT_TRUE(policy.GetAllowKioskAppControlChromeVersion(&bool_value));
-  ASSERT_FALSE(bool_value);
+  EXPECT_FALSE(bool_value);
 
   std::vector<DevicePolicy::UsbDeviceId> list_device;
   ASSERT_TRUE(policy.GetUsbDetachableWhitelist(&list_device));
-  ASSERT_EQ(2, list_device.size());
-  ASSERT_EQ(0x413c, list_device[0].vendor_id);
-  ASSERT_EQ(0x2105, list_device[0].product_id);
-  ASSERT_EQ(0x0403, list_device[1].vendor_id);
-  ASSERT_EQ(0x6001, list_device[1].product_id);
+  EXPECT_EQ(2, list_device.size());
+  EXPECT_EQ(0x413c, list_device[0].vendor_id);
+  EXPECT_EQ(0x2105, list_device[0].product_id);
+  EXPECT_EQ(0x0403, list_device[1].vendor_id);
+  EXPECT_EQ(0x6001, list_device[1].product_id);
 
   ASSERT_TRUE(policy.GetAutoLaunchedKioskAppId(&string_value));
-  ASSERT_EQ("my_kiosk_app", string_value);
+  EXPECT_EQ("my_kiosk_app", string_value);
 
   int_value = -1;
   ASSERT_TRUE(policy.GetSecondFactorAuthenticationMode(&int_value));
-  ASSERT_EQ(2, int_value);
+  EXPECT_EQ(2, int_value);
 
   // Reloading the protobuf should succeed.
-  ASSERT_TRUE(provider.Reload());
+  EXPECT_TRUE(provider.Reload());
 }
 
 // Test that a policy file can be verified and parsed correctly. The file
@@ -185,10 +188,11 @@ TEST(PolicyTest, DevicePolicyNoneSetTest) {
   base::FilePath policy_file(kPolicyFileNoneSet);
   base::FilePath key_file(kKeyFile);
 
-  PolicyProvider provider(
-      CreateDevicePolicyImpl(std::make_unique<MockInstallAttributesReader>(
-                                 cryptohome::SerializedInstallAttributes()),
-                             policy_file, key_file, false));
+  PolicyProvider provider;
+  provider.SetDevicePolicyForTesting(CreateDevicePolicyImpl(
+      std::make_unique<MockInstallAttributesReader>(
+          InstallAttributesReader::kDeviceModeEnterprise, true),
+      policy_file, key_file, false));
   provider.Reload();
 
   // Ensure we successfully loaded the device policy file.
@@ -204,30 +208,33 @@ TEST(PolicyTest, DevicePolicyNoneSetTest) {
   std::string string_value;
   std::vector<DevicePolicy::UsbDeviceId> list_device;
 
-  ASSERT_FALSE(policy.GetPolicyRefreshRate(&int_value));
-  ASSERT_FALSE(policy.GetUserWhitelist(&list_value));
-  ASSERT_FALSE(policy.GetGuestModeEnabled(&bool_value));
-  ASSERT_FALSE(policy.GetCameraEnabled(&bool_value));
-  ASSERT_FALSE(policy.GetShowUserNames(&bool_value));
-  ASSERT_FALSE(policy.GetDataRoamingEnabled(&bool_value));
-  ASSERT_FALSE(policy.GetAllowNewUsers(&bool_value));
-  ASSERT_FALSE(policy.GetMetricsEnabled(&bool_value));
-  ASSERT_FALSE(policy.GetReportVersionInfo(&bool_value));
-  ASSERT_FALSE(policy.GetReportActivityTimes(&bool_value));
-  ASSERT_FALSE(policy.GetReportBootMode(&bool_value));
-  ASSERT_FALSE(policy.GetEphemeralUsersEnabled(&bool_value));
-  ASSERT_FALSE(policy.GetReleaseChannel(&string_value));
-  ASSERT_FALSE(policy.GetUpdateDisabled(&bool_value));
-  ASSERT_FALSE(policy.GetTargetVersionPrefix(&string_value));
-  ASSERT_FALSE(policy.GetRollbackToTargetVersion(&int_value));
-  ASSERT_FALSE(policy.GetRollbackAllowedMilestones(&int_value));
-  ASSERT_FALSE(policy.GetScatterFactorInSeconds(&int64_value));
-  ASSERT_FALSE(policy.GetOpenNetworkConfiguration(&string_value));
-  ASSERT_FALSE(policy.GetHttpDownloadsEnabled(&bool_value));
-  ASSERT_FALSE(policy.GetAuP2PEnabled(&bool_value));
-  ASSERT_FALSE(policy.GetAllowKioskAppControlChromeVersion(&bool_value));
-  ASSERT_FALSE(policy.GetUsbDetachableWhitelist(&list_device));
-  ASSERT_FALSE(policy.GetSecondFactorAuthenticationMode(&int_value));
+  EXPECT_FALSE(policy.GetPolicyRefreshRate(&int_value));
+  EXPECT_FALSE(policy.GetUserWhitelist(&list_value));
+  EXPECT_FALSE(policy.GetGuestModeEnabled(&bool_value));
+  EXPECT_FALSE(policy.GetCameraEnabled(&bool_value));
+  EXPECT_FALSE(policy.GetShowUserNames(&bool_value));
+  EXPECT_FALSE(policy.GetDataRoamingEnabled(&bool_value));
+  EXPECT_FALSE(policy.GetAllowNewUsers(&bool_value));
+  EXPECT_FALSE(policy.GetMetricsEnabled(&bool_value));
+  EXPECT_FALSE(policy.GetReportVersionInfo(&bool_value));
+  EXPECT_FALSE(policy.GetReportActivityTimes(&bool_value));
+  EXPECT_FALSE(policy.GetReportBootMode(&bool_value));
+  EXPECT_FALSE(policy.GetEphemeralUsersEnabled(&bool_value));
+  EXPECT_FALSE(policy.GetReleaseChannel(&string_value));
+  EXPECT_FALSE(policy.GetUpdateDisabled(&bool_value));
+  EXPECT_FALSE(policy.GetTargetVersionPrefix(&string_value));
+  EXPECT_FALSE(policy.GetRollbackToTargetVersion(&int_value));
+  // RollbackAllowedMilestones has the default value of 4 for enterprise
+  // devices.
+  ASSERT_TRUE(policy.GetRollbackAllowedMilestones(&int_value));
+  EXPECT_EQ(4, int_value);
+  EXPECT_FALSE(policy.GetScatterFactorInSeconds(&int64_value));
+  EXPECT_FALSE(policy.GetOpenNetworkConfiguration(&string_value));
+  EXPECT_FALSE(policy.GetHttpDownloadsEnabled(&bool_value));
+  EXPECT_FALSE(policy.GetAuP2PEnabled(&bool_value));
+  EXPECT_FALSE(policy.GetAllowKioskAppControlChromeVersion(&bool_value));
+  EXPECT_FALSE(policy.GetUsbDetachableWhitelist(&list_device));
+  EXPECT_FALSE(policy.GetSecondFactorAuthenticationMode(&int_value));
 }
 
 // Verify that the library will correctly recognize and signal missing files.
@@ -236,14 +243,15 @@ TEST(PolicyTest, DevicePolicyFailure) {
   // Try loading non-existing protobuf should fail.
   base::FilePath policy_file(kNonExistingFile);
   base::FilePath key_file(kNonExistingFile);
-  PolicyProvider provider(
+  PolicyProvider provider;
+  provider.SetDevicePolicyForTesting(
       CreateDevicePolicyImpl(std::make_unique<MockInstallAttributesReader>(
                                  cryptohome::SerializedInstallAttributes()),
                              policy_file, key_file, true));
 
   // Even after reload the policy should still be not loaded.
   ASSERT_FALSE(provider.Reload());
-  ASSERT_FALSE(provider.device_policy_is_loaded());
+  EXPECT_FALSE(provider.device_policy_is_loaded());
 }
 
 // Verify that signature verification is waived for a device in enterprise_ad
@@ -251,20 +259,15 @@ TEST(PolicyTest, DevicePolicyFailure) {
 TEST(PolicyTest, SkipSignatureForEnterpriseAD) {
   base::FilePath policy_file(kPolicyFileAllSet);
   base::FilePath key_file(kNonExistingFile);
-  cryptohome::SerializedInstallAttributes install_attributes;
-  cryptohome::SerializedInstallAttributes::Attribute* attr =
-      install_attributes.add_attributes();
-  ASSERT_NE(nullptr, attr);
-  attr->set_name("enterprise.mode");
-  attr->set_value("enterprise_ad");
-
-  PolicyProvider provider(CreateDevicePolicyImpl(
-      std::make_unique<MockInstallAttributesReader>(install_attributes),
+  PolicyProvider provider;
+  provider.SetDevicePolicyForTesting(CreateDevicePolicyImpl(
+      std::make_unique<MockInstallAttributesReader>(
+          InstallAttributesReader::kDeviceModeEnterpriseAD, true),
       policy_file, key_file, false));
   provider.Reload();
 
   // Ensure we successfully loaded the device policy file.
-  ASSERT_TRUE(provider.device_policy_is_loaded());
+  EXPECT_TRUE(provider.device_policy_is_loaded());
 }
 
 // Ensure that signature verification is enforced for a device in vanilla
@@ -272,20 +275,16 @@ TEST(PolicyTest, SkipSignatureForEnterpriseAD) {
 TEST(PolicyTest, DontSkipSignatureForEnterprise) {
   base::FilePath policy_file(kPolicyFileAllSet);
   base::FilePath key_file(kNonExistingFile);
-  cryptohome::SerializedInstallAttributes install_attributes;
-  cryptohome::SerializedInstallAttributes::Attribute* attr =
-      install_attributes.add_attributes();
-  ASSERT_NE(nullptr, attr);
-  attr->set_name("enterprise.mode");
-  attr->set_value("enterprise");
 
-  PolicyProvider provider(CreateDevicePolicyImpl(
-      std::make_unique<MockInstallAttributesReader>(install_attributes),
+  PolicyProvider provider;
+  provider.SetDevicePolicyForTesting(CreateDevicePolicyImpl(
+      std::make_unique<MockInstallAttributesReader>(
+          InstallAttributesReader::kDeviceModeEnterprise, true),
       policy_file, key_file, false));
   provider.Reload();
 
   // Ensure that unverifed policy is not loaded.
-  ASSERT_FALSE(provider.device_policy_is_loaded());
+  EXPECT_FALSE(provider.device_policy_is_loaded());
 }
 
 // Ensure that signature verification is enforced for a device in consumer mode.
@@ -294,13 +293,48 @@ TEST(PolicyTest, DontSkipSignatureForConsumer) {
   base::FilePath key_file(kNonExistingFile);
   cryptohome::SerializedInstallAttributes install_attributes;
 
-  PolicyProvider provider(CreateDevicePolicyImpl(
+  PolicyProvider provider;
+  provider.SetDevicePolicyForTesting(CreateDevicePolicyImpl(
       std::make_unique<MockInstallAttributesReader>(install_attributes),
       policy_file, key_file, false));
   provider.Reload();
 
   // Ensure that unverifed policy is not loaded.
-  ASSERT_FALSE(provider.device_policy_is_loaded());
+  EXPECT_FALSE(provider.device_policy_is_loaded());
+}
+
+// Checks return value of IsConsumerDevice when it's a still in OOBE.
+TEST(PolicyTest, IsConsumerDeviceOobe) {
+  PolicyProvider provider;
+  provider.SetInstallAttributesReaderForTesting(
+      std::make_unique<MockInstallAttributesReader>("", false));
+  EXPECT_FALSE(provider.IsConsumerDevice());
+}
+
+// Checks return value of IsConsumerDevice when it's a consumer device.
+TEST(PolicyTest, IsConsumerDeviceConsumer) {
+  PolicyProvider provider;
+  provider.SetInstallAttributesReaderForTesting(
+      std::make_unique<MockInstallAttributesReader>("", true));
+  EXPECT_TRUE(provider.IsConsumerDevice());
+}
+
+// Checks return value of IsConsumerDevice when it's an enterprise device.
+TEST(PolicyTest, IsConsumerDeviceEnterprise) {
+  PolicyProvider provider;
+  provider.SetInstallAttributesReaderForTesting(
+      std::make_unique<MockInstallAttributesReader>(
+          InstallAttributesReader::kDeviceModeEnterprise, true));
+  EXPECT_FALSE(provider.IsConsumerDevice());
+}
+
+// Checks return value of IsConsumerDevice when it's an enterprise AD device.
+TEST(PolicyTest, IsConsumerDeviceEnterpriseAd) {
+  PolicyProvider provider;
+  provider.SetInstallAttributesReaderForTesting(
+      std::make_unique<MockInstallAttributesReader>(
+          InstallAttributesReader::kDeviceModeEnterpriseAD, true));
+  EXPECT_FALSE(provider.IsConsumerDevice());
 }
 
 }  // namespace policy