Reland "DeviceStateMixin: Update local state and support for setting policy"

This is a reland of d80f755b

TBR=antrim@chromium.org
TBR=alemate@chromium.org

Original change's description:
> DeviceStateMixin: Update local state and support for setting policy
>
> Taking over rsorokin's patch at CL:1578742, given that he is OOO.
>
> DeviceStateMixin now initialized device state values in local
> state during setup.
>
> Add support for setting device policy and device local account policies
> in in-memory fake session manager client, in a way similar to
> UserPolicyMixin. Later this functionality can be moved to a separate
> mixin, which might make sense if/when we add functionality for setting
> up policy served by local test policy server to this mixin (this should
> already be possible using LocalPolicyTestServerMixin).
>
> Uses AutoLaunchedKioskTest as an example of new usage.
>
> Also fixes a number of tests to wirk with the new device state setup.
>
> Bug: 952855
> Change-Id: I1c0314a51581137d0f5a0907a09a4ef5b0bccda3
> Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1600857
> Reviewed-by: Denis Kuznetsov <antrim@chromium.org>
> Reviewed-by: Alexander Alekseev <alemate@chromium.org>
> Commit-Queue: Toni Baržić <tbarzic@chromium.org>
> Cr-Commit-Position: refs/heads/master@{#658297}

Bug: 952855
Change-Id: I966406b55aaffe78137be296a31b4b42202fe8cd
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1604186Reviewed-by: Toni Baržić <tbarzic@chromium.org>
Reviewed-by: Denis Kuznetsov <antrim@chromium.org>
Commit-Queue: Toni Baržić <tbarzic@chromium.org>
Cr-Commit-Position: refs/heads/master@{#658808}

chrome/browser/chromeos/login/auto_launched_kiosk_browsertest.cc

@@ -7,14 +7,10 @@
 #include <vector>
 
 #include "apps/test/app_window_waiter.h"
-#include "base/base64.h"
 #include "base/command_line.h"
 #include "base/files/file_path.h"
-#include "base/json/json_file_value_serializer.h"
 #include "base/logging.h"
-#include "base/path_service.h"
 #include "base/run_loop.h"
-#include "base/stl_util.h"
 #include "base/values.h"
 #include "chrome/browser/chrome_notification_types.h"
 #include "chrome/browser/chromeos/app_mode/fake_cws.h"
@@ -22,17 +18,13 @@
 #include "chrome/browser/chromeos/app_mode/kiosk_app_manager.h"
 #include "chrome/browser/chromeos/login/app_launch_controller.h"
 #include "chrome/browser/chromeos/login/mixin_based_in_process_browser_test.h"
+#include "chrome/browser/chromeos/login/test/device_state_mixin.h"
 #include "chrome/browser/chromeos/login/test/embedded_test_server_mixin.h"
 #include "chrome/browser/chromeos/login/test/login_manager_mixin.h"
-#include "chrome/browser/chromeos/ownership/owner_settings_service_chromeos_factory.h"
 #include "chrome/browser/chromeos/policy/device_local_account.h"
-#include "chrome/browser/chromeos/policy/device_policy_builder.h"
-#include "chrome/browser/chromeos/policy/device_policy_cros_browser_test.h"
 #include "chrome/browser/extensions/browsertest_util.h"
-#include "chrome/browser/extensions/extension_apitest.h"
 #include "chrome/browser/profiles/profile_manager.h"
 #include "chrome/common/chrome_constants.h"
-#include "chrome/common/chrome_paths.h"
 #include "chrome/common/pref_names.h"
 #include "chromeos/dbus/dbus_thread_manager.h"
 #include "chromeos/dbus/session_manager/fake_session_manager_client.h"
@@ -46,7 +38,7 @@
 #include "extensions/browser/app_window/app_window.h"
 #include "extensions/browser/app_window/app_window_registry.h"
 #include "extensions/browser/app_window/native_app_window.h"
-#include "extensions/common/value_builder.h"
+#include "extensions/browser/sandboxed_unpacker.h"
 #include "extensions/test/extension_test_message_listener.h"
 #include "extensions/test/result_catcher.h"
 #include "net/dns/mock_host_resolver.h"
@@ -127,10 +119,9 @@ class TerminationObserver : public content::NotificationObserver {
 class AutoLaunchedKioskTest : public MixinBasedInProcessBrowserTest {
  public:
   AutoLaunchedKioskTest()
-      : install_attributes_(
-            chromeos::StubInstallAttributes::CreateCloudManaged("domain.com",
-                                                                "device_id")),
-        verifier_format_override_(crx_file::VerifierFormat::CRX3) {}
+      : verifier_format_override_(crx_file::VerifierFormat::CRX3) {
+    device_state_.set_domain("domain.com");
+  }
 
   ~AutoLaunchedKioskTest() override = default;
 
@@ -159,34 +150,30 @@ class AutoLaunchedKioskTest : public MixinBasedInProcessBrowserTest {
     MixinBasedInProcessBrowserTest::SetUpCommandLine(command_line);
   }
 
-  bool SetUpUserDataDirectory() override {
-    InitDevicePolicy();
-
-    base::FilePath user_data_path;
-    if (!base::PathService::Get(chrome::DIR_USER_DATA, &user_data_path)) {
-      ADD_FAILURE() << "Unable to get used data dir";
-      return false;
-    }
-
-    if (!CacheDevicePolicyToLocalState(user_data_path))
-      return false;
-
-    return MixinBasedInProcessBrowserTest::SetUpUserDataDirectory();
-  }
-
   void SetUpInProcessBrowserTestFixture() override {
     host_resolver()->AddRule("*", "127.0.0.1");
 
     SessionManagerClient::InitializeFakeInMemory();
 
     FakeSessionManagerClient::Get()->set_supports_browser_restart(true);
-    FakeSessionManagerClient::Get()->set_device_policy(
-        device_policy_helper_.device_policy()->GetBlob());
-    FakeSessionManagerClient::Get()->set_device_local_account_policy(
-        kTestAccountId, device_local_account_policy_.GetBlob());
 
-    // Arbitrary non-empty state keys.
-    FakeSessionManagerClient::Get()->set_server_backed_state_keys({"1"});
+    std::unique_ptr<ScopedDevicePolicyUpdate> device_policy_update =
+        device_state_.RequestDevicePolicyUpdate();
+    em::DeviceLocalAccountsProto* const device_local_accounts =
+        device_policy_update->policy_payload()->mutable_device_local_accounts();
+    device_local_accounts->set_auto_login_id(kTestAccountId);
+
+    em::DeviceLocalAccountInfoProto* const account =
+        device_local_accounts->add_account();
+    account->set_account_id(kTestAccountId);
+    account->set_type(em::DeviceLocalAccountInfoProto::ACCOUNT_TYPE_KIOSK_APP);
+    account->mutable_kiosk_app()->set_app_id(GetTestAppId());
+
+    device_policy_update.reset();
+
+    std::unique_ptr<ScopedUserPolicyUpdate> device_local_account_policy_update =
+        device_state_.RequestDeviceLocalAccountPolicyUpdate(kTestAccountId);
+    device_local_account_policy_update.reset();
 
     MixinBasedInProcessBrowserTest::SetUpInProcessBrowserTestFixture();
   }
@@ -215,61 +202,6 @@ class AutoLaunchedKioskTest : public MixinBasedInProcessBrowserTest {
     MixinBasedInProcessBrowserTest::TearDownOnMainThread();
   }
 
-  void InitDevicePolicy() {
-    device_policy_helper_.InstallOwnerKey();
-
-    // Create device policy, and cache it to local state.
-    em::DeviceLocalAccountsProto* const device_local_accounts =
-        device_policy_helper_.device_policy()
-            ->payload()
-            .mutable_device_local_accounts();
-
-    em::DeviceLocalAccountInfoProto* const account =
-        device_local_accounts->add_account();
-    account->set_account_id(kTestAccountId);
-    account->set_type(em::DeviceLocalAccountInfoProto::ACCOUNT_TYPE_KIOSK_APP);
-    account->mutable_kiosk_app()->set_app_id(GetTestAppId());
-
-    device_local_accounts->set_auto_login_id(kTestAccountId);
-
-    device_policy_helper_.device_policy()->Build();
-
-    device_local_account_policy_.policy_data().set_username(kTestAccountId);
-    device_local_account_policy_.policy_data().set_policy_type(
-        policy::dm_protocol::kChromePublicAccountPolicyType);
-    device_local_account_policy_.policy_data().set_settings_entity_id(
-        kTestAccountId);
-    device_local_account_policy_.Build();
-  }
-
-  bool CacheDevicePolicyToLocalState(const base::FilePath& user_data_path) {
-    em::PolicyData policy_data;
-    if (!device_policy_helper_.device_policy()->payload().SerializeToString(
-            policy_data.mutable_policy_value())) {
-      ADD_FAILURE() << "Failed to serialize device policy.";
-      return false;
-    }
-    const std::string policy_data_str = policy_data.SerializeAsString();
-    std::string policy_data_encoded;
-    base::Base64Encode(policy_data_str, &policy_data_encoded);
-
-    std::unique_ptr<base::DictionaryValue> local_state =
-        extensions::DictionaryBuilder()
-            .Set(prefs::kDeviceSettingsCache, policy_data_encoded)
-            .Set("PublicAccounts",
-                 extensions::ListBuilder().Append(GetTestAppUserId()).Build())
-            .Build();
-    local_state->SetKey(prefs::kOobeComplete, base::Value(true));
-
-    JSONFileValueSerializer serializer(
-        user_data_path.Append(chrome::kLocalStateFilename));
-    if (!serializer.Serialize(*local_state)) {
-      ADD_FAILURE() << "Failed to write local state.";
-      return false;
-    }
-    return true;
-  }
-
   const std::string GetTestAppUserId() const {
     return policy::GenerateDeviceLocalAccountUserId(
         kTestAccountId, policy::DeviceLocalAccount::TYPE_KIOSK_APP);
@@ -322,9 +254,6 @@ class AutoLaunchedKioskTest : public MixinBasedInProcessBrowserTest {
   std::unique_ptr<TerminationObserver> termination_observer_;
 
  private:
-  chromeos::ScopedStubInstallAttributes install_attributes_;
-  policy::UserPolicyBuilder device_local_account_policy_;
-  policy::DevicePolicyCrosTestHelper device_policy_helper_;
   FakeCWS fake_cws_;
   extensions::SandboxedUnpacker::ScopedVerifierFormatOverrideForTest
       verifier_format_override_;
@@ -333,6 +262,9 @@ class AutoLaunchedKioskTest : public MixinBasedInProcessBrowserTest {
       &mixin_host_, embedded_test_server()};
   LoginManagerMixin login_manager_{&mixin_host_, {}};
 
+  DeviceStateMixin device_state_{
+      &mixin_host_, DeviceStateMixin::State::OOBE_COMPLETED_CLOUD_ENROLLED};
+
   DISALLOW_COPY_AND_ASSIGN(AutoLaunchedKioskTest);
 };
 

chrome/browser/chromeos/login/encryption_migration_browsertest.cc

@@ -92,7 +92,7 @@ class EncryptionMigrationTest : public MixinBasedInProcessBrowserTest {
 
   void SetUpEncryptionMigrationActionPolicy(
       arc::policy_util::EcryptfsMigrationAction action) {
-    std::unique_ptr<UserPolicyMixin::ScopedPolicyUpdate> updater =
+    std::unique_ptr<ScopedUserPolicyUpdate> updater =
         user_policy_mixin_.RequestCachedPolicyUpdate();
     updater->policy_payload()->mutable_ecryptfsmigrationstrategy()->set_value(
         static_cast<int>(action));

chrome/browser/chromeos/login/enrollment/enrollment_local_policy_server_browsertest.cc

@@ -185,7 +185,7 @@ class AutoEnrollmentNoStateKeys : public AutoEnrollmentWithStatistics {
   // AutoEnrollmentWithStatistics:
   void SetUpInProcessBrowserTestFixture() override {
     AutoEnrollmentWithStatistics::SetUpInProcessBrowserTestFixture();
-    SessionManagerClient::InitializeFakeInMemory();
+    // Session manager client is initialized by DeviceStateMixin.
     FakeSessionManagerClient::Get()->set_force_state_keys_missing(true);
   }
 

chrome/browser/chromeos/login/existing_user_controller_browsertest.cc

@@ -221,6 +221,16 @@ class ExistingUserControllerTest : public policy::DevicePolicyCrosBrowserTest {
     users_pref->AppendIfNotPresent(std::make_unique<base::Value>(user_id));
   }
 
+  void MakeCrosSettingsPermanentlyUntrusted() {
+    device_policy()->policy().set_policy_data_signature("bad signature");
+    session_manager_client()->set_device_policy(device_policy()->GetBlob());
+    session_manager_client()->OnPropertyChangeComplete(true);
+
+    base::RunLoop run_loop;
+    WaitForPermanentlyUntrustedStatusAndRun(run_loop.QuitClosure());
+    run_loop.Run();
+  }
+
   // ExistingUserController private member accessors.
   base::OneShotTimer* auto_login_timer() {
     return existing_user_controller()->auto_login_timer_.get();
@@ -288,8 +298,9 @@ class ExistingUserControllerUntrustedTest : public ExistingUserControllerTest {
  public:
   ExistingUserControllerUntrustedTest() = default;
 
-  void SetUpInProcessBrowserTestFixture() override {
-    ExistingUserControllerTest::SetUpInProcessBrowserTestFixture();
+  void SetUpOnMainThread() override {
+    ExistingUserControllerTest::SetUpOnMainThread();
+    MakeCrosSettingsPermanentlyUntrusted();
     ExpectLoginFailure();
   }
 
@@ -493,16 +504,6 @@ class ExistingUserControllerPublicSessionTest
     existing_user_controller()->OnPublicSessionAutoLoginTimerFire();
   }
 
-  void MakeCrosSettingsPermanentlyUntrusted() {
-    device_policy()->policy().set_policy_data_signature("bad signature");
-    session_manager_client()->set_device_policy(device_policy()->GetBlob());
-    session_manager_client()->OnPropertyChangeComplete(true);
-
-    base::RunLoop run_loop;
-    WaitForPermanentlyUntrustedStatusAndRun(run_loop.QuitClosure());
-    run_loop.Run();
-  }
-
   const AccountId public_session_account_id_ =
       AccountId::FromUserEmail(policy::GenerateDeviceLocalAccountUserId(
           kPublicSessionUserEmail,

chrome/browser/chromeos/login/saml/saml_browsertest.cc

@@ -38,6 +38,7 @@
 #include "chrome/browser/chromeos/login/test/https_forwarder.h"
 #include "chrome/browser/chromeos/login/test/js_checker.h"
 #include "chrome/browser/chromeos/login/test/local_policy_test_server_mixin.h"
+#include "chrome/browser/chromeos/login/test/login_manager_mixin.h"
 #include "chrome/browser/chromeos/login/test/login_screen_tester.h"
 #include "chrome/browser/chromeos/login/test/oobe_base_test.h"
 #include "chrome/browser/chromeos/login/test/oobe_screen_waiter.h"
@@ -932,6 +933,14 @@ class SAMLPolicyTest : public SamlTest {
       &mixin_host_,
       chromeos::DeviceStateMixin::State::OOBE_COMPLETED_CLOUD_ENROLLED};
 
+  // Add a fake user so the login screen does not show GAIA auth by default.
+  // This enables tests to control when the GAIA is shown (and ensure it's
+  // loaded after SAML config has been set up).
+  chromeos::LoginManagerMixin login_manager_{
+      &mixin_host_,
+      {chromeos::LoginManagerMixin::TestUserInfo(
+          AccountId::FromUserEmailGaiaId("user@gmail.com", "1111"))}};
+
  private:
   DISALLOW_COPY_AND_ASSIGN(SAMLPolicyTest);
 };
@@ -952,7 +961,9 @@ class SAMLPolicyTestWebUILogin : public SAMLPolicyTest {
 };
 
 SAMLPolicyTest::SAMLPolicyTest()
-    : device_policy_(test_helper_.device_policy()) {}
+    : device_policy_(test_helper_.device_policy()) {
+  device_state_.set_skip_initial_policy_setup(true);
+}
 
 SAMLPolicyTest::~SAMLPolicyTest() {}
 
@@ -1090,6 +1101,8 @@ void SAMLPolicyTest::ShowGAIALoginForm() {
 
 void SAMLPolicyTest::ShowSAMLInterstitial() {
   login_screen_load_observer_->Wait();
+  WaitForOobeUI();
+
   content::DOMMessageQueue message_queue;
   ASSERT_TRUE(
       content::ExecuteScript(GetLoginUI()->GetWebContents(),

chrome/browser/chromeos/login/test/device_state_mixin.cc

@@ -4,12 +4,23 @@
 
 #include "chrome/browser/chromeos/login/test/device_state_mixin.h"
 
+#include <utility>
+#include <vector>
+
+#include "base/callback.h"
+#include "base/json/json_writer.h"
 #include "base/numerics/safe_conversions.h"
 #include "base/path_service.h"
 #include "chrome/browser/browser_process.h"
+#include "chrome/browser/chrome_browser_main.h"
+#include "chrome/browser/chrome_browser_main_extra_parts.h"
+#include "chrome/browser/chromeos/policy/device_policy_builder.h"
 #include "chrome/common/chrome_paths.h"
 #include "chrome/common/pref_names.h"
 #include "chromeos/constants/chromeos_paths.h"
+#include "chromeos/dbus/session_manager/fake_session_manager_client.h"
+#include "chromeos/dbus/session_manager/session_manager_client.h"
+#include "components/policy/core/common/cloud/cloud_policy_constants.h"
 #include "components/policy/core/common/cloud/policy_builder.h"
 #include "components/policy/proto/install_attributes.pb.h"
 #include "components/prefs/pref_service.h"
@@ -48,6 +59,55 @@ cryptohome::SerializedInstallAttributes BuildInstallAttributes(
   return install_attrs;
 }
 
+void WriteFile(const base::FilePath& path, const std::string& blob) {
+  CHECK_EQ(base::checked_cast<int>(blob.length()),
+           base::WriteFile(path, blob.data(), blob.length()));
+}
+
+// Set of values that should be written to local state.
+struct LocalStateValues {
+  base::Optional<bool> oobe_complete;
+  base::Optional<bool> enrollment_recovery_required;
+  base::Optional<int> device_registered;
+};
+
+// Chrome main extra part used to initialize local state prefs to indicate the
+// configured device state. Injected into browser main parts in
+// CreatedBrowserMainParts() override.
+class TestLocalStateInitializerMainExtra : public ChromeBrowserMainExtraParts {
+ public:
+  explicit TestLocalStateInitializerMainExtra(
+      LocalStateValues local_state_values)
+      : local_state_values_(std::move(local_state_values)) {}
+  ~TestLocalStateInitializerMainExtra() override = default;
+
+  // ChromeBrowserMainExtraParts:
+  void PostEarlyInitialization() override {
+    PrefService* local_state = g_browser_process->local_state();
+
+    if (local_state_values_.oobe_complete.has_value()) {
+      local_state->SetBoolean(prefs::kOobeComplete,
+                              *local_state_values_.oobe_complete);
+    }
+
+    if (local_state_values_.enrollment_recovery_required.has_value()) {
+      local_state->SetBoolean(
+          prefs::kEnrollmentRecoveryRequired,
+          *local_state_values_.enrollment_recovery_required);
+    }
+
+    if (local_state_values_.device_registered.has_value()) {
+      local_state->SetInteger(prefs::kDeviceRegistered,
+                              *local_state_values_.device_registered);
+    }
+  }
+
+ private:
+  LocalStateValues local_state_values_;
+
+  DISALLOW_COPY_AND_ASSIGN(TestLocalStateInitializerMainExtra);
+};
+
 }  // namespace
 
 DeviceStateMixin::DeviceStateMixin(InProcessBrowserTestMixinHost* host,
@@ -62,6 +122,75 @@ bool DeviceStateMixin::SetUpUserDataDirectory() {
   return true;
 }
 
+void DeviceStateMixin::SetUpInProcessBrowserTestFixture() {
+  // Make sure session manager client has been initialized as in-memory. This is
+  // requirement for setting policy blobs.
+  if (!chromeos::SessionManagerClient::Get())
+    chromeos::SessionManagerClient::InitializeFakeInMemory();
+
+  session_manager_initialized_ = true;
+
+  std::vector<std::string> state_keys;
+  state_keys.push_back("1");
+  FakeSessionManagerClient::Get()->set_server_backed_state_keys(state_keys);
+
+  if (IsEnrolledState() && !skip_initial_policy_setup_) {
+    SetCachedDevicePolicy();
+
+    for (const auto& device_local_account : device_local_account_policies_)
+      SetCachedDeviceLocalAccountPolicy(device_local_account.first);
+  }
+}
+
+void DeviceStateMixin::CreatedBrowserMainParts(
+    content::BrowserMainParts* browser_main_parts) {
+  LocalStateValues local_state_values;
+
+  switch (state_) {
+    case DeviceStateMixin::State::OOBE_COMPLETED_CLOUD_ENROLLED:
+    case DeviceStateMixin::State::OOBE_COMPLETED_ACTIVE_DIRECTORY_ENROLLED:
+    case DeviceStateMixin::State::OOBE_COMPLETED_CONSUMER_OWNED:
+    case DeviceStateMixin::State::OOBE_COMPLETED_DEMO_MODE:
+      local_state_values.device_registered = 1;
+      FALLTHROUGH;
+    case DeviceStateMixin::State::OOBE_COMPLETED_UNOWNED:
+      local_state_values.enrollment_recovery_required = false;
+      local_state_values.oobe_complete = true;
+      break;
+    case DeviceStateMixin::State::BEFORE_OOBE:
+      break;
+  }
+
+  // |browser_main_parts| take ownership of TestUserRegistrationMainExtra.
+  static_cast<ChromeBrowserMainParts*>(browser_main_parts)
+      ->AddParts(new TestLocalStateInitializerMainExtra(
+          std::move(local_state_values)));
+}
+
+std::unique_ptr<ScopedDevicePolicyUpdate>
+DeviceStateMixin::RequestDevicePolicyUpdate() {
+  if (!IsEnrolledState())
+    return nullptr;
+
+  return std::make_unique<ScopedDevicePolicyUpdate>(
+      &device_policy_, base::BindOnce(&DeviceStateMixin::SetCachedDevicePolicy,
+                                      weak_factory_.GetWeakPtr()));
+}
+
+std::unique_ptr<ScopedUserPolicyUpdate>
+DeviceStateMixin::RequestDeviceLocalAccountPolicyUpdate(
+    const std::string& account_id) {
+  if (!IsEnrolledState())
+    return nullptr;
+
+  policy::UserPolicyBuilder& builder =
+      device_local_account_policies_[account_id];
+  return std::make_unique<ScopedUserPolicyUpdate>(
+      &builder,
+      base::BindRepeating(&DeviceStateMixin::SetCachedDeviceLocalAccountPolicy,
+                          weak_factory_.GetWeakPtr(), account_id));
+}
+
 void DeviceStateMixin::SetState(State state) {
   DCHECK(!is_setup_) << "SetState called after device was set up";
   state_ = state;
@@ -110,11 +239,10 @@ void DeviceStateMixin::WriteInstallAttrFile() {
       break;
   }
 
-  std::string blob;
+  std::string install_attrs_bits;
   CHECK(BuildInstallAttributes(device_mode, domain, realm, kFakeDeviceId)
-            .SerializeToString(&blob));
-  CHECK_EQ(base::checked_cast<int>(blob.length()),
-           base::WriteFile(install_attrs_file, blob.data(), blob.length()));
+            .SerializeToString(&install_attrs_bits));
+  WriteFile(install_attrs_file, install_attrs_bits);
 }
 
 void DeviceStateMixin::WriteOwnerKey() {
@@ -126,18 +254,63 @@ void DeviceStateMixin::WriteOwnerKey() {
     case DeviceStateMixin::State::OOBE_COMPLETED_CLOUD_ENROLLED:
     case DeviceStateMixin::State::OOBE_COMPLETED_CONSUMER_OWNED:
     case DeviceStateMixin::State::OOBE_COMPLETED_DEMO_MODE:
-      base::FilePath user_data_dir;
-      base::FilePath owner_key_file;
-      CHECK(base::PathService::Get(chrome::DIR_USER_DATA, &user_data_dir));
-      owner_key_file = user_data_dir.Append("stub_owner.key");
-      const std::string owner_key_bits =
-          policy::PolicyBuilder::GetPublicTestKeyAsString();
-      CHECK(!owner_key_bits.empty());
-      CHECK_EQ(base::checked_cast<int>(owner_key_bits.length()),
-               base::WriteFile(owner_key_file, owner_key_bits.data(),
-                               owner_key_bits.length()));
       break;
   }
+
+  base::FilePath user_data_dir;
+  CHECK(base::PathService::Get(chrome::DIR_USER_DATA, &user_data_dir));
+  base::FilePath owner_key_file = user_data_dir.Append("stub_owner.key");
+  const std::string owner_key_bits =
+      policy::PolicyBuilder::GetPublicTestKeyAsString();
+  CHECK(!owner_key_bits.empty());
+  WriteFile(owner_key_file, owner_key_bits);
+}
+
+bool DeviceStateMixin::IsEnrolledState() const {
+  switch (state_) {
+    case DeviceStateMixin::State::BEFORE_OOBE:
+    case DeviceStateMixin::State::OOBE_COMPLETED_UNOWNED:
+    case DeviceStateMixin::State::OOBE_COMPLETED_CONSUMER_OWNED:
+      return false;
+    case DeviceStateMixin::State::OOBE_COMPLETED_ACTIVE_DIRECTORY_ENROLLED:
+    case DeviceStateMixin::State::OOBE_COMPLETED_CLOUD_ENROLLED:
+    case DeviceStateMixin::State::OOBE_COMPLETED_DEMO_MODE:
+      return true;
+  }
+  return false;
+}
+
+void DeviceStateMixin::SetCachedDevicePolicy() {
+  if (!session_manager_initialized_)
+    return;
+
+  DCHECK(IsEnrolledState());
+
+  device_policy_.SetDefaultSigningKey();
+  device_policy_.Build();
+  FakeSessionManagerClient::Get()->set_device_policy(device_policy_.GetBlob());
+  FakeSessionManagerClient::Get()->OnPropertyChangeComplete(true);
+}
+
+void DeviceStateMixin::SetCachedDeviceLocalAccountPolicy(
+    const std::string& account_id) {
+  if (!session_manager_initialized_ ||
+      !device_local_account_policies_.count(account_id))
+    return;
+
+  DCHECK(IsEnrolledState());
+
+  policy::UserPolicyBuilder& builder =
+      device_local_account_policies_[account_id];
+  builder.policy_data().set_username(account_id);
+  builder.policy_data().set_settings_entity_id(account_id);
+  builder.policy_data().set_policy_type(
+      policy::dm_protocol::kChromePublicAccountPolicyType);
+  builder.SetDefaultSigningKey();
+  builder.Build();
+
+  FakeSessionManagerClient::Get()->set_device_local_account_policy(
+      account_id, builder.GetBlob());
 }
 
 DeviceStateMixin::~DeviceStateMixin() = default;

chrome/browser/chromeos/login/test/device_state_mixin.h

@@ -5,16 +5,31 @@
 #ifndef CHROME_BROWSER_CHROMEOS_LOGIN_TEST_DEVICE_STATE_MIXIN_H_
 #define CHROME_BROWSER_CHROMEOS_LOGIN_TEST_DEVICE_STATE_MIXIN_H_
 
+#include <map>
 #include <memory>
 #include <string>
 
 #include "base/macros.h"
+#include "base/memory/weak_ptr.h"
 #include "base/values.h"
 #include "chrome/browser/chromeos/login/demo_mode/demo_session.h"
 #include "chrome/browser/chromeos/login/mixin_based_in_process_browser_test.h"
+#include "chrome/browser/chromeos/login/test/scoped_policy_update.h"
 
 namespace chromeos {
 
+// A mixin for setting up device state:
+// *   OOBE completion state
+// *   enrollment state
+// *   install attributes
+// *   fake owner key
+// *   cached device and local account policies
+//
+// Note that this uses in memory fake session manager client to store the
+// policy blobs.
+// It will initialized the fake in-memory client in
+// SetUpInProcessBrowserTestFixture(), provided that a session manager was not
+// initialized previously.
 class DeviceStateMixin : public InProcessBrowserTestMixin {
  public:
   enum class State {
@@ -31,20 +46,75 @@ class DeviceStateMixin : public InProcessBrowserTestMixin {
 
   // InProcessBrowserTestMixin:
   bool SetUpUserDataDirectory() override;
+  void SetUpInProcessBrowserTestFixture() override;
+  void CreatedBrowserMainParts(
+      content::BrowserMainParts* browser_main_parts) override;
+
+  // Returns a ScopedDevicePolicyUpdate instance that can be used to update
+  // local device policy blob (kept in fake session manager client).
+  // The cached policy value will be updated as the returned
+  // ScopedDevicePolicyUpdate goes out of scope or when
+  // FakeSessionManagerClient is initialized in
+  // SetUpInProcessBrowserTestFixture(), whichever happens later.
+  //
+  // NOTE: If your test is serving policy using test policy server, use this to
+  // set initially cached policy blob only. Prefre using policy server for
+  // policy updates.
+  std::unique_ptr<ScopedDevicePolicyUpdate> RequestDevicePolicyUpdate();
+
+  // Returns ScopedUserPolicyUpdate instance that can be used to set, or update
+  // local account policy blob for |account_id| (kept in fake session manager
+  // client).
+  //
+  // The cached policy value will be updated as the returned
+  // ScopedUsetPolicyUpdate goes out of scope or FakeSessionManagerClient is
+  // initialized in SetUpInProcessBrowserTestFixtire(), whichever happens later.
+  // DeviceStateMixin will fill out the required policy data values.
+  //
+  // NOTE: If your test is serving policy using test policy server, use this to
+  // set initially cached policy blob only. Prefre using policy server for
+  // policy updates.
+  std::unique_ptr<ScopedUserPolicyUpdate> RequestDeviceLocalAccountPolicyUpdate(
+      const std::string& account_id);
 
   void SetState(State state);
   void set_domain(const std::string& domain) { domain_ = domain; }
+  void set_skip_initial_policy_setup(bool value) {
+    skip_initial_policy_setup_ = value;
+  }
 
  private:
   void SetDeviceState();
   void WriteInstallAttrFile();
   void WriteOwnerKey();
 
+  // Whether |state_| value indicates enrolled state.
+  bool IsEnrolledState() const;
+
+  // Updates device policy blob stored by fake session manager client.
+  void SetCachedDevicePolicy();
+
+  // Updates device local account policy blob stored by fake session manager
+  // client.
+  void SetCachedDeviceLocalAccountPolicy(const std::string& account_id);
+
   State state_;
   std::string domain_;
 
   bool is_setup_ = false;
 
+  // If set, the mixin will not set any device policy values during test setup,
+  // unless requested using RequestDevicePolicyUpdate().
+  bool skip_initial_policy_setup_ = false;
+
+  // Whether session manager client has been initialized.
+  bool session_manager_initialized_ = false;
+  policy::DevicePolicyBuilder device_policy_;
+  std::map<std::string, policy::UserPolicyBuilder>
+      device_local_account_policies_;
+
+  base::WeakPtrFactory<DeviceStateMixin> weak_factory_{this};
+
   DISALLOW_COPY_AND_ASSIGN(DeviceStateMixin);
 };
 

chrome/browser/chromeos/login/test/oobe_base_test.cc

@@ -62,13 +62,20 @@ void OobeBaseTest::SetUpCommandLine(base::CommandLine* command_line) {
   MixinBasedInProcessBrowserTest::SetUpCommandLine(command_line);
 }
 
-void OobeBaseTest::SetUpOnMainThread() {
-  host_resolver()->AddRule("*", "127.0.0.1");
-
+void OobeBaseTest::CreatedBrowserMainParts(
+    content::BrowserMainParts* browser_main_parts) {
+  // If the test initially shows views login screen, this notification might
+  // come before SetUpOnMainThread(), so the observer has to be set up early.
   login_screen_load_observer_.reset(new content::WindowedNotificationObserver(
       chrome::NOTIFICATION_LOGIN_OR_LOCK_WEBUI_VISIBLE,
       content::NotificationService::AllSources()));
 
+  MixinBasedInProcessBrowserTest::CreatedBrowserMainParts(browser_main_parts);
+}
+
+void OobeBaseTest::SetUpOnMainThread() {
+  host_resolver()->AddRule("*", "127.0.0.1");
+
   test::UserSessionManagerTestApi session_manager_test_api(
       UserSessionManager::GetInstance());
   session_manager_test_api.SetShouldObtainTokenHandleInTests(false);

chrome/browser/chromeos/login/test/oobe_base_test.h

@@ -5,6 +5,7 @@
 #ifndef CHROME_BROWSER_CHROMEOS_LOGIN_TEST_OOBE_BASE_TEST_H_
 #define CHROME_BROWSER_CHROMEOS_LOGIN_TEST_OOBE_BASE_TEST_H_
 
+#include <memory>
 #include <string>
 
 #include "base/macros.h"
@@ -33,6 +34,8 @@ class OobeBaseTest : public MixinBasedInProcessBrowserTest {
   // MixinBasedInProcessBrowserTest::
   void SetUp() override;
   void SetUpCommandLine(base::CommandLine* command_line) override;
+  void CreatedBrowserMainParts(
+      content::BrowserMainParts* browser_main_parts) override;
   void SetUpOnMainThread() override;
 
   // If this returns true (default), then SetUpOnMainThread would wait for

chrome/browser/chromeos/login/test/scoped_policy_update.cc

+// Copyright 2019 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/login/test/scoped_policy_update.h"
+
+#include <utility>
+
+namespace chromeos {
+
+ScopedUserPolicyUpdate::ScopedUserPolicyUpdate(
+    policy::UserPolicyBuilder* policy_builder,
+    base::OnceClosure callback)
+    : policy_builder_(policy_builder), callback_(std::move(callback)) {}
+
+ScopedUserPolicyUpdate::~ScopedUserPolicyUpdate() {
+  std::move(callback_).Run();
+}
+
+ScopedDevicePolicyUpdate::ScopedDevicePolicyUpdate(
+    policy::DevicePolicyBuilder* policy_builder,
+    base::OnceClosure callback)
+    : policy_builder_(policy_builder), callback_(std::move(callback)) {}
+
+ScopedDevicePolicyUpdate::~ScopedDevicePolicyUpdate() {
+  std::move(callback_).Run();
+}
+
+}  // namespace chromeos

chrome/browser/chromeos/login/test/scoped_policy_update.h

+// Copyright 2019 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_LOGIN_TEST_SCOPED_POLICY_UPDATE_H_
+#define CHROME_BROWSER_CHROMEOS_LOGIN_TEST_SCOPED_POLICY_UPDATE_H_
+
+#include "base/callback.h"
+#include "base/macros.h"
+#include "chrome/browser/chromeos/policy/device_policy_builder.h"
+#include "components/policy/core/common/cloud/policy_builder.h"
+
+namespace chromeos {
+
+// Helper that is used by mixins to provide means for setting up user policy
+// values to tests that need that functionality. This does not build, nor apply
+// updated policy - that's expected to be done in |callback|.
+//
+// |callback| - Called when this object goes out of scope.
+class ScopedUserPolicyUpdate {
+ public:
+  explicit ScopedUserPolicyUpdate(policy::UserPolicyBuilder* policy_builder,
+                                  base::OnceClosure callback);
+  ~ScopedUserPolicyUpdate();
+
+  // Policy payload proto - use this to set up desired policy values.
+  enterprise_management::CloudPolicySettings* policy_payload() {
+    return &policy_builder_->payload();
+  }
+
+ private:
+  policy::UserPolicyBuilder* const policy_builder_;
+  base::OnceClosure callback_;
+
+  DISALLOW_COPY_AND_ASSIGN(ScopedUserPolicyUpdate);
+};
+
+// Helper that is used by mixins to provide means for setting up device policy
+// values to tests that need that functionality. This does not build, nor apply
+// updated policy - that's expected to be done in |callback|.
+//
+// |callback| - Called when this object goes out of scope.
+class ScopedDevicePolicyUpdate {
+ public:
+  explicit ScopedDevicePolicyUpdate(policy::DevicePolicyBuilder* policy_builder,
+                                    base::OnceClosure callback);
+  ~ScopedDevicePolicyUpdate();
+
+  // Policy payload proto - use this to set up desired policy values.
+  enterprise_management::ChromeDeviceSettingsProto* policy_payload() {
+    return &policy_builder_->payload();
+  }
+
+  // The policy data that will be used to build the policy.
+  enterprise_management::PolicyData* policy_data() {
+    return &policy_builder_->policy_data();
+  }
+
+ private:
+  policy::DevicePolicyBuilder* const policy_builder_;
+  base::OnceClosure callback_;
+
+  DISALLOW_COPY_AND_ASSIGN(ScopedDevicePolicyUpdate);
+};
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_CHROMEOS_LOGIN_TEST_SCOPED_POLICY_UPDATE_H_

chrome/browser/chromeos/login/test/user_policy_mixin.cc

@@ -21,15 +21,6 @@
 
 namespace chromeos {
 
-UserPolicyMixin::ScopedPolicyUpdate::ScopedPolicyUpdate(
-    policy::UserPolicyBuilder* policy_builder,
-    base::OnceClosure callback)
-    : policy_builder_(policy_builder), callback_(std::move(callback)) {}
-
-UserPolicyMixin::ScopedPolicyUpdate::~ScopedPolicyUpdate() {
-  std::move(callback_).Run();
-}
-
 UserPolicyMixin::UserPolicyMixin(InProcessBrowserTestMixinHost* mixin_host,
                                  const AccountId& account_id)
     : InProcessBrowserTestMixin(mixin_host), account_id_(account_id) {}
@@ -50,9 +41,9 @@ void UserPolicyMixin::SetUpInProcessBrowserTestFixture() {
     SetUpCachedPolicy();
 }
 
-std::unique_ptr<UserPolicyMixin::ScopedPolicyUpdate>
+std::unique_ptr<ScopedUserPolicyUpdate>
 UserPolicyMixin::RequestCachedPolicyUpdate() {
-  return std::make_unique<ScopedPolicyUpdate>(
+  return std::make_unique<ScopedUserPolicyUpdate>(
       &cached_user_policy_builder_,
       base::BindOnce(&UserPolicyMixin::SetUpCachedPolicy,
                      weak_factory_.GetWeakPtr()));

chrome/browser/chromeos/login/test/user_policy_mixin.h

@@ -8,10 +8,10 @@
 #include <memory>
 #include <string>
 
-#include "base/callback.h"
 #include "base/macros.h"
 #include "base/memory/weak_ptr.h"
 #include "chrome/browser/chromeos/login/mixin_based_in_process_browser_test.h"
+#include "chrome/browser/chromeos/login/test/scoped_policy_update.h"
 #include "components/account_id/account_id.h"
 #include "components/policy/core/common/cloud/policy_builder.h"
 
@@ -22,28 +22,6 @@ namespace chromeos {
 // NOTE: This mixin will set up in-memory FakeSessionManagerClient during setup.
 class UserPolicyMixin : public InProcessBrowserTestMixin {
  public:
-  // Helper that provides access to policy payload to test that need to update
-  // it,
-  // |callback| - callback that applies policy changes. Called when this object
-  // goes out of scope.
-  class ScopedPolicyUpdate {
-   public:
-    explicit ScopedPolicyUpdate(policy::UserPolicyBuilder* policy_builder,
-                                base::OnceClosure callback);
-    ~ScopedPolicyUpdate();
-
-    // Policy payload proto - use this to set up desired policy values.
-    enterprise_management::CloudPolicySettings* policy_payload() {
-      return &policy_builder_->payload();
-    }
-
-   private:
-    policy::UserPolicyBuilder* const policy_builder_;
-    base::OnceClosure callback_;
-
-    DISALLOW_COPY_AND_ASSIGN(ScopedPolicyUpdate);
-  };
-
   UserPolicyMixin(InProcessBrowserTestMixinHost* mixin_host,
                   const AccountId& account_id);
   ~UserPolicyMixin() override;
@@ -51,14 +29,14 @@ class UserPolicyMixin : public InProcessBrowserTestMixin {
   // InProcessBrowserTestMixin:
   void SetUpInProcessBrowserTestFixture() override;
 
-  // Returns a ScopedPolicyUpdate object that will update the cached policy
+  // Returns a ScopedUserPolicyUpdate object that will update the cached policy
   // values as it goes out of scope. Calling this will ensure that the cached
   // policy blob is set (even if policy remains empty when ScopedPolicyUpdate is
   // done).
   //
   // If called during setup, before steps that initialize session manager,
   // policy change will be deferred until session manager initialization.
-  std::unique_ptr<ScopedPolicyUpdate> RequestCachedPolicyUpdate();
+  std::unique_ptr<ScopedUserPolicyUpdate> RequestCachedPolicyUpdate();
 
  private:
   // Creates a file containing public policy signing key that will be used to

chrome/browser/chromeos/login/webview_login_browsertest.cc

@@ -441,17 +441,16 @@ class WebviewClientCertsLoginTest : public WebviewLoginTest {
   // Sets up the DeviceLoginScreenAutoSelectCertificateForUrls policy.
   void SetAutoSelectCertificatePatterns(
       const std::vector<std::string>& autoselect_patterns) {
-    em::ChromeDeviceSettingsProto& proto(
-        device_policy_test_helper_.device_policy()->payload());
+    em::ChromeDeviceSettingsProto& proto(device_policy_builder_.payload());
     auto* field =
         proto.mutable_device_login_screen_auto_select_certificate_for_urls();
     for (const std::string& autoselect_pattern : autoselect_patterns)
       field->add_login_screen_auto_select_certificate_rules(autoselect_pattern);
 
-    device_policy_test_helper_.device_policy()->Build();
+    device_policy_builder_.Build();
 
     FakeSessionManagerClient::Get()->set_device_policy(
-        device_policy_test_helper_.device_policy()->GetBlob());
+        device_policy_builder_.GetBlob());
     PrefChangeWatcher watcher(prefs::kManagedAutoSelectCertificateForUrls,
                               ProfileHelper::GetSigninProfile()->GetPrefs());
     FakeSessionManagerClient::Get()->OnPropertyChangeComplete(true);
@@ -471,16 +470,15 @@ class WebviewClientCertsLoginTest : public WebviewLoginTest {
     base::DictionaryValue onc_dict =
         BuildDeviceOncDictForUntrustedAuthority(x509_contents);
 
-    em::ChromeDeviceSettingsProto& proto(
-        device_policy_test_helper_.device_policy()->payload());
+    em::ChromeDeviceSettingsProto& proto(device_policy_builder_.payload());
     base::JSONWriter::Write(onc_dict,
                             proto.mutable_open_network_configuration()
                                 ->mutable_open_network_configuration());
 
-    device_policy_test_helper_.device_policy()->Build();
+    device_policy_builder_.Build();
 
     FakeSessionManagerClient::Get()->set_device_policy(
-        device_policy_test_helper_.device_policy()->GetBlob());
+        device_policy_builder_.GetBlob());
     PrefChangeWatcher watcher(onc::prefs::kDeviceOpenNetworkConfiguration,
                               g_browser_process->local_state());
     FakeSessionManagerClient::Get()->OnPropertyChangeComplete(true);
@@ -545,13 +543,17 @@ class WebviewClientCertsLoginTest : public WebviewLoginTest {
   void SetUpInProcessBrowserTestFixture() override {
     // Override FakeSessionManagerClient. This will be shut down by the browser.
     chromeos::SessionManagerClient::InitializeFakeInMemory();
+    device_policy_builder_.Build();
     FakeSessionManagerClient::Get()->set_device_policy(
-        device_policy_test_helper_.device_policy()->GetBlob());
+        device_policy_builder_.GetBlob());
 
     WebviewLoginTest::SetUpInProcessBrowserTestFixture();
   }
 
-  void TearDownOnMainThread() override { TearDownTestSystemSlot(); }
+  void TearDownOnMainThread() override {
+    TearDownTestSystemSlot();
+    WebviewLoginTest::TearDownOnMainThread();
+  }
 
  private:
   void SetUpTestSystemSlotOnIO(bool* out_system_slot_constructed_successfully) {
@@ -599,7 +601,7 @@ class WebviewClientCertsLoginTest : public WebviewLoginTest {
     return onc_dict;
   }
 
-  policy::DevicePolicyCrosTestHelper device_policy_test_helper_;
+  policy::DevicePolicyBuilder device_policy_builder_;
   std::unique_ptr<crypto::ScopedTestSystemNSSKeySlot> test_system_slot_;
   scoped_refptr<net::X509Certificate> client_cert_;
   std::unique_ptr<net::SpawnedTestServer> https_server_;
@@ -879,10 +881,8 @@ class WebviewProxyAuthLoginTest : public WebviewLoginTest {
     //   registered for policy.
     // - the payload is given to |policy_test_server_|, so we can download fresh
     //   policy.
-    device_policy_test_helper_.device_policy()
-        ->policy_data()
-        .set_public_key_version(1);
-    device_policy_test_helper_.device_policy()->Build();
+    device_policy_builder()->policy_data().set_public_key_version(1);
+    device_policy_builder()->Build();
 
     UpdateServedPolicyFromDevicePolicyTestHelper();
     WebviewLoginTest::SetUp();
@@ -958,7 +958,7 @@ class WebviewProxyAuthLoginTest : public WebviewLoginTest {
   }
 
   policy::DevicePolicyBuilder* device_policy_builder() {
-    return device_policy_test_helper_.device_policy();
+    return &device_policy_builder_;
   }
 
   content::WindowedNotificationObserver* auth_needed_observer() {
@@ -975,7 +975,7 @@ class WebviewProxyAuthLoginTest : public WebviewLoginTest {
   // authentication method.
   std::unique_ptr<net::SpawnedTestServer> auth_proxy_server_;
   LocalPolicyTestServerMixin local_policy_mixin_{&mixin_host_};
-  policy::DevicePolicyCrosTestHelper device_policy_test_helper_;
+  policy::DevicePolicyBuilder device_policy_builder_;
 
   DeviceStateMixin device_state_{
       &mixin_host_, DeviceStateMixin::State::OOBE_COMPLETED_CLOUD_ENROLLED};

chrome/browser/chromeos/policy/browser_policy_connector_chromeos_browsertest.cc

@@ -41,6 +41,11 @@ void WaitUntilPolicyLoaded() {
 }
 
 class BrowserPolicyConnectorChromeOSTest : public DevicePolicyCrosBrowserTest {
+ public:
+  BrowserPolicyConnectorChromeOSTest() {
+    device_state_.set_skip_initial_policy_setup(true);
+  }
+  ~BrowserPolicyConnectorChromeOSTest() override = default;
 };
 
 // Test that GetEnterpriseEnrollmentDomain and GetEnterpriseDisplayDomain work

chrome/browser/chromeos/policy/device_local_account_browsertest.cc

@@ -13,6 +13,7 @@
 #include <utility>
 #include <vector>
 
+#include "ash/public/cpp/ash_switches.h"
 #include "ash/shell.h"
 #include "ash/system/session/logout_confirmation_controller.h"
 #include "ash/system/session/logout_confirmation_dialog.h"
@@ -49,6 +50,7 @@
 #include "chrome/browser/chromeos/login/session/user_session_manager_test_api.h"
 #include "chrome/browser/chromeos/login/signin_specifics.h"
 #include "chrome/browser/chromeos/login/test/local_policy_test_server_mixin.h"
+#include "chrome/browser/chromeos/login/test/oobe_screen_waiter.h"
 #include "chrome/browser/chromeos/login/ui/login_display_host.h"
 #include "chrome/browser/chromeos/login/ui/webui_login_view.h"
 #include "chrome/browser/chromeos/login/users/avatar/user_image_manager.h"
@@ -84,6 +86,7 @@
 #include "chrome/browser/ui/extensions/app_launch_params.h"
 #include "chrome/browser/ui/extensions/application_launch.h"
 #include "chrome/browser/ui/tabs/tab_strip_model.h"
+#include "chrome/browser/ui/webui/chromeos/login/gaia_screen_handler.h"
 #include "chrome/browser/ui/webui/chromeos/login/oobe_ui.h"
 #include "chrome/browser/ui/webui/chromeos/login/signin_screen_handler.h"
 #include "chrome/browser/ui/webui/chromeos/login/terms_of_service_screen_handler.h"
@@ -469,6 +472,10 @@ class DeviceLocalAccountTest : public DevicePolicyCrosBrowserTest,
     command_line->AppendSwitch(chromeos::switches::kLoginManager);
     command_line->AppendSwitch(chromeos::switches::kForceLoginManagerInTests);
     command_line->AppendSwitchASCII(chromeos::switches::kLoginProfile, "user");
+
+    // LoginDisplayHostMojo does not support dynamic device policy changes (see
+    // https://crbug.com/956456).
+    command_line->AppendSwitch(ash::switches::kShowWebUiLogin);
   }
 
   void SetUpInProcessBrowserTestFixture() override {
@@ -509,10 +516,7 @@ class DeviceLocalAccountTest : public DevicePolicyCrosBrowserTest,
       run_loop.Run();
 
     // Skip to the login screen.
-    chromeos::WizardController* wizard_controller =
-        chromeos::WizardController::default_controller();
-    ASSERT_TRUE(wizard_controller);
-    wizard_controller->SkipToLoginForTesting(LoginScreenContext());
+    chromeos::OobeScreenWaiter(chromeos::GaiaView::kScreenId).Wait();
 
     chromeos::test::UserSessionManagerTestApi session_manager_test_api(
         chromeos::UserSessionManager::GetInstance());

chrome/browser/chromeos/policy/device_policy_cros_browser_test.cc

@@ -63,14 +63,6 @@ DevicePolicyCrosBrowserTest::DevicePolicyCrosBrowserTest()
 
 DevicePolicyCrosBrowserTest::~DevicePolicyCrosBrowserTest() = default;
 
-void DevicePolicyCrosBrowserTest::SetUp() {
-  // Set some fake state keys to make surethey are not empty.
-  std::vector<std::string> state_keys;
-  state_keys.push_back("1");
-  fake_session_manager_client_->set_server_backed_state_keys(state_keys);
-  chromeos::MixinBasedInProcessBrowserTest::SetUp();
-}
-
 void DevicePolicyCrosBrowserTest::RefreshDevicePolicy() {
   // Reset the key to its original state.
   device_policy()->SetDefaultSigningKey();

chrome/browser/chromeos/policy/device_policy_cros_browser_test.h

@@ -46,9 +46,6 @@ class DevicePolicyCrosBrowserTest
   DevicePolicyCrosBrowserTest();
   ~DevicePolicyCrosBrowserTest() override;
 
-  // MixinBasedInProcessBrowserTest:
-  void SetUp() override;
-
   // Reinstalls |device_policy_| as the policy (to be used when it was
   // recently changed).
   void RefreshDevicePolicy();

chrome/browser/chromeos/policy/display_resolution_handler_browsertest.cc

@@ -382,7 +382,9 @@ class DisplayResolutionBootTest
     : public chromeos::MixinBasedInProcessBrowserTest,
       public testing::WithParamInterface<PolicyValue> {
  protected:
-  DisplayResolutionBootTest() = default;
+  DisplayResolutionBootTest() {
+    device_state_.set_skip_initial_policy_setup(true);
+  }
   ~DisplayResolutionBootTest() override = default;
 
   void SetUpCommandLine(base::CommandLine* command_line) override {

chrome/browser/chromeos/policy/display_rotation_default_handler_browsertest.cc

@@ -259,7 +259,9 @@ class DisplayRotationBootTest
     : public chromeos::MixinBasedInProcessBrowserTest,
       public testing::WithParamInterface<display::Display::Rotation> {
  protected:
-  DisplayRotationBootTest() = default;
+  DisplayRotationBootTest() {
+    device_state_.set_skip_initial_policy_setup(true);
+  }
   ~DisplayRotationBootTest() override = default;
 
   void SetUpInProcessBrowserTestFixture() override {

chrome/browser/chromeos/policy/login_policy_test_base.cc

@@ -97,14 +97,7 @@ void LoginPolicyTestBase::SetMergeSessionParams() {
 
 void LoginPolicyTestBase::SkipToLoginScreen() {
   chromeos::WizardController::SkipPostLoginScreensForTesting();
-  chromeos::WizardController* const wizard_controller =
-      chromeos::WizardController::default_controller();
-  ASSERT_TRUE(wizard_controller);
-  wizard_controller->SkipToLoginForTesting(chromeos::LoginScreenContext());
-
-  content::WindowedNotificationObserver(
-      chrome::NOTIFICATION_LOGIN_OR_LOCK_WEBUI_VISIBLE,
-      content::NotificationService::AllSources()).Wait();
+  OobeBaseTest::WaitForSigninScreen();
 }
 
 void LoginPolicyTestBase::LogIn(const std::string& user_id,

chrome/browser/chromeos/policy/site_isolation_flag_handling_browsertest.cc

@@ -16,8 +16,10 @@
 #include "chrome/browser/chromeos/login/session/user_session_manager.h"
 #include "chrome/browser/chromeos/login/session/user_session_manager_test_api.h"
 #include "chrome/browser/chromeos/login/test/device_state_mixin.h"
+#include "chrome/browser/chromeos/login/test/user_policy_mixin.h"
 #include "chrome/browser/chromeos/login/ui/login_display_host.h"
 #include "chrome/browser/chromeos/login/users/chrome_user_manager_impl.h"
+#include "chrome/browser/chromeos/login/wizard_controller.h"
 #include "chrome/browser/chromeos/policy/browser_policy_connector_chromeos.h"
 #include "chrome/browser/chromeos/policy/login_policy_test_base.h"
 #include "chrome/browser/chromeos/settings/cros_settings.h"
@@ -254,16 +256,18 @@ const Params kTestCases[] = {
            false /* expected_request_restart */,
            {} /* expected_flags_for_user */)};
 
+constexpr char kTestUserAccountId[] = "username@examle.com";
 constexpr char kTestUserGaiaId[] = "1111111111";
+constexpr char kTestUserPassword[] = "password";
 
 class SiteIsolationFlagHandlingTest
-    : public policy::LoginPolicyTestBase,
+    : public OobeBaseTest,
       public ::testing::WithParamInterface<Params> {
  protected:
   SiteIsolationFlagHandlingTest() = default;
 
   void SetUpCommandLine(base::CommandLine* command_line) override {
-    policy::LoginPolicyTestBase::SetUpCommandLine(command_line);
+    OobeBaseTest::SetUpCommandLine(command_line);
 
     // Simulate login_manager behavior: pass --site-per-process or
     // --isolate-origins between policy flag sentinels according to test case
@@ -285,39 +289,51 @@ class SiteIsolationFlagHandlingTest
       command_line->AppendSwitch(switches::kPolicySwitchesEnd);
   }
 
-  // This is called from |LoginPolicyTestBase| and specifies the user policy
-  // values to be served by the local policy test server.
-  void GetMandatoryPoliciesValue(base::DictionaryValue* policy) const override {
-    if (GetParam().user_policy_site_per_process) {
-      policy->Set(policy::key::kSitePerProcess,
-                  std::make_unique<base::Value>(true));
-    }
-    if (!GetParam().user_policy_isolate_origins.empty()) {
-      policy->Set(policy::key::kIsolateOrigins,
-                  std::make_unique<base::Value>(
-                      GetParam().user_policy_isolate_origins));
-    }
-  }
 
   void SetUpInProcessBrowserTestFixture() override {
-    policy::LoginPolicyTestBase::SetUpInProcessBrowserTestFixture();
+    SessionManagerClient::InitializeFakeInMemory();
 
     // Mark that chrome restart can be requested.
     // Note that AttemptRestart() is mocked out in UserSessionManager through
     // |SetAttemptRestartClosureInTests| (set up in SetUpOnMainThread).
-    SessionManagerClient::InitializeFake();
     FakeSessionManagerClient::Get()->set_supports_browser_restart(true);
+
+    std::unique_ptr<ScopedDevicePolicyUpdate> update =
+        device_state_.RequestDevicePolicyUpdate();
+    update->policy_payload()
+        ->mutable_ephemeral_users_enabled()
+        ->set_ephemeral_users_enabled(GetParam().ephemeral_users);
+    update.reset();
+
+    std::unique_ptr<ScopedUserPolicyUpdate> user_policy_update =
+        user_policy_.RequestCachedPolicyUpdate();
+    if (GetParam().user_policy_site_per_process) {
+      user_policy_update->policy_payload()
+          ->mutable_siteperprocess()
+          ->mutable_policy_options()
+          ->set_mode(em::PolicyOptions::MANDATORY);
+      user_policy_update->policy_payload()->mutable_siteperprocess()->set_value(
+          true);
+    }
+
+    if (!GetParam().user_policy_isolate_origins.empty()) {
+      user_policy_update->policy_payload()
+          ->mutable_isolateorigins()
+          ->mutable_policy_options()
+          ->set_mode(em::PolicyOptions::MANDATORY);
+      user_policy_update->policy_payload()->mutable_isolateorigins()->set_value(
+          GetParam().user_policy_isolate_origins);
+    }
+    user_policy_update.reset();
+
+    OobeBaseTest::SetUpInProcessBrowserTestFixture();
   }
 
   void SetUpOnMainThread() override {
-    policy::LoginPolicyTestBase::SetUpOnMainThread();
-
-    // Write ephemeral users status directly into CrosSettings.
-    scoped_testing_cros_settings_.device_settings()->SetBoolean(
-        kAccountsPrefEphemeralUsersEnabled, GetParam().ephemeral_users);
+    fake_gaia_.SetupFakeGaiaForLogin(kTestUserAccountId, kTestUserGaiaId,
+                                     "fake-refresh-token");
 
-    // This makes the user manager reload CrosSettings.
-    GetChromeUserManager()->OwnershipStatusChanged();
+    OobeBaseTest::SetUpOnMainThread();
 
     login_wait_loop_ = std::make_unique<base::RunLoop>();
 
@@ -369,12 +385,15 @@ class SiteIsolationFlagHandlingTest
   // This is important because ephemeral users only work on enrolled machines.
   DeviceStateMixin device_state_{
       &mixin_host_, DeviceStateMixin::State::OOBE_COMPLETED_CLOUD_ENROLLED};
+  UserPolicyMixin user_policy_{
+      &mixin_host_,
+      AccountId::FromUserEmailGaiaId(kTestUserAccountId, kTestUserGaiaId)};
+
+  chromeos::FakeGaiaMixin fake_gaia_{&mixin_host_, embedded_test_server()};
 
   // Observes for user session start.
   std::unique_ptr<content::WindowedNotificationObserver>
       user_session_started_observer_;
-  policy::MockConfigurationPolicyProvider provider_;
-  chromeos::ScopedTestingCrosSettings scoped_testing_cros_settings_;
   DISALLOW_COPY_AND_ASSIGN(SiteIsolationFlagHandlingTest);
 };
 
@@ -384,11 +403,13 @@ IN_PROC_BROWSER_TEST_P(SiteIsolationFlagHandlingTest, FlagHandlingTest) {
   // Start user sign-in. We can't use |LoginPolicyTestBase::LogIn|, because
   // it waits for a user session start unconditionally, which will not happen if
   // chrome requests a restart to set user-session flags.
-  SkipToLoginScreen();
+  chromeos::WizardController::SkipPostLoginScreensForTesting();
+  OobeBaseTest::WaitForSigninScreen();
+
   LoginDisplayHost::default_host()
       ->GetOobeUI()
       ->GetView<GaiaScreenHandler>()
-      ->ShowSigninScreenForTest(kAccountId, kAccountPassword, kEmptyServices);
+      ->ShowSigninScreenForTest(kTestUserAccountId, kTestUserPassword, "[]");
 
   // Wait for either the user session to start, or for restart to be requested
   // (whichever happens first).
@@ -408,7 +429,7 @@ IN_PROC_BROWSER_TEST_P(SiteIsolationFlagHandlingTest, FlagHandlingTest) {
 
   // Also verify flags if chrome was restarted.
   AccountId test_account_id =
-      AccountId::FromUserEmailGaiaId(GetAccount(), kTestUserGaiaId);
+      AccountId::FromUserEmailGaiaId(kTestUserAccountId, kTestUserGaiaId);
   std::vector<std::string> flags_for_user;
   bool has_flags_for_user = FakeSessionManagerClient::Get()->GetFlagsForUser(
       cryptohome::CreateAccountIdentifierFromAccountId(test_account_id),

chrome/browser/chromeos/system/device_disabling_browsertest.cc

@@ -17,8 +17,6 @@
 #include "chrome/browser/chromeos/login/ui/login_display_host.h"
 #include "chrome/browser/chromeos/login/ui/webui_login_view.h"
 #include "chrome/browser/chromeos/login/wizard_controller.h"
-#include "chrome/browser/chromeos/policy/device_policy_builder.h"
-#include "chrome/browser/chromeos/policy/device_policy_cros_browser_test.h"
 #include "chrome/browser/chromeos/profiles/profile_helper.h"
 #include "chrome/browser/chromeos/settings/cros_settings.h"
 #include "chrome/browser/ui/webui/chromeos/login/device_disabled_screen_handler.h"
@@ -75,7 +73,6 @@ class DeviceDisablingTest
 
  protected:
   // OobeBaseTest:
-  void SetUpInProcessBrowserTestFixture() override;
   void SetUpOnMainThread() override;
 
   // NetworkStateInformer::NetworkStateInformerObserver:
@@ -85,8 +82,6 @@ class DeviceDisablingTest
   NetworkPortalDetectorMixin network_portal_detector_{&mixin_host_};
 
  private:
-  policy::DevicePolicyCrosTestHelper test_helper_;
-
   chromeos::DeviceStateMixin device_state_{
       &mixin_host_,
       chromeos::DeviceStateMixin::State::OOBE_COMPLETED_CLOUD_ENROLLED};
@@ -97,11 +92,10 @@ class DeviceDisablingTest
 
 void DeviceDisablingTest::SetDeviceDisabledPolicy() {
   // Prepare a policy fetch response that indicates the device is disabled.
-  test_helper_.device_policy()->policy_data().mutable_device_state()->
-      set_device_mode(enterprise_management::DeviceState::DEVICE_MODE_DISABLED);
-  test_helper_.device_policy()->Build();
-  FakeSessionManagerClient::Get()->set_device_policy(
-      test_helper_.device_policy()->GetBlob());
+  std::unique_ptr<ScopedDevicePolicyUpdate> policy_update =
+      device_state_.RequestDevicePolicyUpdate();
+  policy_update->policy_data()->mutable_device_state()->set_device_mode(
+      enterprise_management::DeviceState::DEVICE_MODE_DISABLED);
 }
 
 void DeviceDisablingTest::MarkDisabledAndWaitForPolicyFetch() {
@@ -129,13 +123,6 @@ std::string DeviceDisablingTest::GetCurrentScreenName(
   return screen_name;
 }
 
-void DeviceDisablingTest::SetUpInProcessBrowserTestFixture() {
-  // Override FakeSessionManagerClient. This will be shut down by the browser.
-  chromeos::SessionManagerClient::InitializeFakeInMemory();
-
-  OobeBaseTest::SetUpInProcessBrowserTestFixture();
-}
-
 void DeviceDisablingTest::SetUpOnMainThread() {
   network_state_change_wait_run_loop_.reset(new base::RunLoop);
 
@@ -172,9 +159,6 @@ IN_PROC_BROWSER_TEST_F(DeviceDisablingTest, DisableWithEphemeralUsers) {
   connect_run_loop.Run();
 
   // Skip to the login screen.
-  WizardController* wizard_controller = WizardController::default_controller();
-  ASSERT_TRUE(wizard_controller);
-  wizard_controller->SkipToLoginForTesting(LoginScreenContext());
   OobeScreenWaiter(GaiaView::kScreenId).Wait();
 
   // Mark the device as disabled and wait until cros settings update.

chrome/test/BUILD.gn

@@ -1935,6 +1935,8 @@ test("browser_tests") {
         "../browser/chromeos/login/test/oobe_base_test.h",
         "../browser/chromeos/login/test/oobe_window_visibility_waiter.cc",
         "../browser/chromeos/login/test/oobe_window_visibility_waiter.h",
+        "../browser/chromeos/login/test/scoped_policy_update.cc",
+        "../browser/chromeos/login/test/scoped_policy_update.h",
         "../browser/chromeos/login/test/session_flags_manager.cc",
         "../browser/chromeos/login/test/session_flags_manager.h",
         "../browser/chromeos/login/test/session_manager_state_waiter.cc",