Add license selection test based on local policy server

Bug: 950471
Change-Id: Ic3ae509848884bb8349543fda157edbb35c70cee
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1556801
Commit-Queue: Denis Kuznetsov <antrim@chromium.org>
Reviewed-by: Roman Sorokin [CET] <rsorokin@chromium.org>
Cr-Commit-Position: refs/heads/master@{#648653}

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

@@ -8,6 +8,7 @@
 #include "chrome/browser/chromeos/login/enrollment/auto_enrollment_check_screen.h"
 #include "chrome/browser/chromeos/login/enrollment/enrollment_screen.h"
 #include "chrome/browser/chromeos/login/startup_utils.h"
+#include "chrome/browser/chromeos/login/test/enrollment_ui_mixin.h"
 #include "chrome/browser/chromeos/login/test/fake_gaia_mixin.h"
 #include "chrome/browser/chromeos/login/test/js_checker.h"
 #include "chrome/browser/chromeos/login/test/local_policy_test_server_mixin.h"
@@ -82,6 +83,7 @@ class EnrollmentLocalPolicyServerBase : public OobeBaseTest {
   }
 
   LocalPolicyTestServerMixin local_policy_mixin_{&mixin_host_};
+  test::EnrollmentUIMixin enrollment_ui_{&mixin_host_};
   FakeGaiaMixin fake_gaia_{&mixin_host_, embedded_test_server()};
 
  private:
@@ -134,8 +136,55 @@ IN_PROC_BROWSER_TEST_F(EnrollmentLocalPolicyServerBase, ManualEnrollment) {
   enrollment_screen()->OnLoginDone(FakeGaiaMixin::kFakeUserEmail,
                                    FakeGaiaMixin::kFakeAuthCode);
 
-  OobeBaseTest::WaitForEnrollmentSuccess();
-  // TODO(rsorokin): Interact with attribute prompt step.
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepDeviceAttributes);
+  enrollment_ui_.SubmitDeviceAttributes(test::values::kAssetId,
+                                        test::values::kLocation);
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepSuccess);
+  EXPECT_TRUE(StartupUtils::IsDeviceRegistered());
+}
+
+// Simple manual enrollment with only license type available.
+// Client should automatically select the only available license type,
+// so no license selection UI should be displayed.
+IN_PROC_BROWSER_TEST_F(EnrollmentLocalPolicyServerBase,
+                       ManualEnrollmentWithSingleLicense) {
+  local_policy_mixin_.ExpectAvailableLicenseCount(
+      5 /* perpetual */, 0 /* annual */, 0 /* kiosk */);
+  host()->StartWizard(OobeScreen::SCREEN_OOBE_ENROLLMENT);
+  OobeScreenWaiter(OobeScreen::SCREEN_OOBE_ENROLLMENT).Wait();
+
+  ASSERT_FALSE(StartupUtils::IsDeviceRegistered());
+  enrollment_screen()->OnLoginDone(FakeGaiaMixin::kFakeUserEmail,
+                                   FakeGaiaMixin::kFakeAuthCode);
+
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepDeviceAttributes);
+  enrollment_ui_.SubmitDeviceAttributes(test::values::kAssetId,
+                                        test::values::kLocation);
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepSuccess);
+  EXPECT_TRUE(StartupUtils::IsDeviceRegistered());
+}
+
+// Simple manual enrollment with license selection.
+// Enrollment selection UI should be displayed during enrollment.
+IN_PROC_BROWSER_TEST_F(EnrollmentLocalPolicyServerBase,
+                       ManualEnrollmentWithMultipleLicenses) {
+  local_policy_mixin_.ExpectAvailableLicenseCount(
+      5 /* perpetual */, 5 /* annual */, 5 /* kiosk */);
+  host()->StartWizard(OobeScreen::SCREEN_OOBE_ENROLLMENT);
+  OobeScreenWaiter(OobeScreen::SCREEN_OOBE_ENROLLMENT).Wait();
+
+  ASSERT_FALSE(StartupUtils::IsDeviceRegistered());
+  enrollment_screen()->OnLoginDone(FakeGaiaMixin::kFakeUserEmail,
+                                   FakeGaiaMixin::kFakeAuthCode);
+
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepLicenses);
+  enrollment_ui_.SelectEnrollmentLicense(test::values::kLicenseTypeAnnual);
+  enrollment_ui_.UseSelectedLicense();
+
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepDeviceAttributes);
+  enrollment_ui_.SubmitDeviceAttributes(test::values::kAssetId,
+                                        test::values::kLocation);
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepSuccess);
   EXPECT_TRUE(StartupUtils::IsDeviceRegistered());
 }
 
@@ -205,7 +254,7 @@ IN_PROC_BROWSER_TEST_F(AutoEnrollmentLocalPolicyServer, Attestation) {
       kTestDomain));
 
   host()->StartWizard(OobeScreen::SCREEN_AUTO_ENROLLMENT_CHECK);
-  OobeBaseTest::WaitForEnrollmentSuccess();
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepDeviceAttributes);
   EXPECT_TRUE(StartupUtils::IsDeviceRegistered());
 }
 

chrome/browser/chromeos/login/enterprise_enrollment_browsertest.cc

@@ -16,6 +16,7 @@
 #include "chrome/browser/chromeos/login/login_manager_test.h"
 #include "chrome/browser/chromeos/login/startup_utils.h"
 #include "chrome/browser/chromeos/login/test/enrollment_helper_mixin.h"
+#include "chrome/browser/chromeos/login/test/enrollment_ui_mixin.h"
 #include "chrome/browser/chromeos/login/test/js_checker.h"
 #include "chrome/browser/chromeos/login/test/oobe_screen_waiter.h"
 #include "chrome/browser/chromeos/login/ui/login_display_host.h"
@@ -147,13 +148,6 @@ class EnterpriseEnrollmentTestBase : public LoginManagerTest {
     ExecutePendingJavaScript();
   }
 
-  // Fills out the UI with device attribute information and submits it.
-  void SubmitAttributePromptUpdate() {
-    // Fill out the attribute prompt info and submit it.
-    test::OobeJS().TypeIntoPath("asset_id", {"oauth-enroll-asset-id"});
-    test::OobeJS().TypeIntoPath("location", {"oauth-enroll-location"});
-    test::OobeJS().TapOn("enroll-attributes-submit-button");
-  }
 
   // Completes the enrollment process.
   void CompleteEnrollment() {
@@ -170,23 +164,6 @@ class EnterpriseEnrollmentTestBase : public LoginManagerTest {
   // and the test code.
   void ExecutePendingJavaScript() { test::OobeJS().Evaluate(";"); }
 
-  // Returns true if there are any DOM elements with the given class.
-  bool IsStepDisplayed(const std::string& step) {
-    const std::string js =
-        "document.getElementsByClassName('oauth-enroll-state-" + step +
-        "').length";
-    int count = test::OobeJS().GetInt(js);
-    return count > 0;
-  }
-
-  // Waits until specific enrollment step is displayed.
-  void WaitForStep(const std::string& step) {
-    const std::string js =
-        "document.getElementsByClassName('oauth-enroll-state-" + step +
-        "').length > 0";
-    test::OobeJS().CreateWaiter(js)->Wait();
-  }
-
   // Setup the enrollment screen.
   void ShowEnrollmentScreen() {
     LoginDisplayHost* host = LoginDisplayHost::default_host();
@@ -205,6 +182,7 @@ class EnterpriseEnrollmentTestBase : public LoginManagerTest {
   }
 
  protected:
+  test::EnrollmentUIMixin enrollment_ui_{&mixin_host_};
   test::EnrollmentHelperMixin enrollment_helper_{&mixin_host_};
 
  private:
@@ -240,7 +218,8 @@ class ActiveDirectoryJoinTest : public EnterpriseEnrollmentTest {
   }
 
   void CheckActiveDirectoryCredentialsShown() {
-    EXPECT_TRUE(IsStepDisplayed("ad-join"));
+    EXPECT_TRUE(
+        enrollment_ui_.IsStepDisplayed(test::ui::kEnrollmentStepADJoin));
     test::OobeJS().ExpectVisiblePath({kAdDialog, kAdCredentialsStep});
     test::OobeJS().ExpectHiddenPath({kAdDialog, kAdUnlockConfigurationStep});
   }
@@ -253,7 +232,8 @@ class ActiveDirectoryJoinTest : public EnterpriseEnrollmentTest {
   }
 
   void CheckActiveDirectoryUnlockConfigurationShown() {
-    EXPECT_TRUE(IsStepDisplayed("ad-join"));
+    EXPECT_TRUE(
+        enrollment_ui_.IsStepDisplayed(test::ui::kEnrollmentStepADJoin));
     test::OobeJS().ExpectHiddenPath({kAdDialog, kAdCredentialsStep});
     test::OobeJS().ExpectVisiblePath({kAdDialog, kAdUnlockConfigurationStep});
   }
@@ -456,8 +436,7 @@ IN_PROC_BROWSER_TEST_F(EnterpriseEnrollmentTest,
   ExecutePendingJavaScript();
 
   // Verify that the error page is displayed.
-  WaitForStep("error");
-  EXPECT_FALSE(IsStepDisplayed("success"));
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepError);
 }
 
 // Shows the enrollment screen and simulates a successful enrollment. Verifies
@@ -472,8 +451,7 @@ IN_PROC_BROWSER_TEST_F(EnterpriseEnrollmentTest,
   SubmitEnrollmentCredentials();
 
   // Verify that the success page is displayed.
-  WaitForStep("success");
-  EXPECT_FALSE(IsStepDisplayed("error"));
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepSuccess);
 }
 
 // Shows the enrollment screen and mocks the enrollment helper to request an
@@ -485,19 +463,17 @@ IN_PROC_BROWSER_TEST_F(EnterpriseEnrollmentTest,
   ShowEnrollmentScreen();
   enrollment_helper_.ExpectEnrollmentMode(
       policy::EnrollmentConfig::MODE_MANUAL);
-  enrollment_helper_.ExpectAttributePromptUpdate("asset_id", "location");
+  enrollment_helper_.ExpectAttributePromptUpdate(test::values::kAssetId,
+                                                 test::values::kLocation);
   enrollment_helper_.ExpectSuccessfulOAuthEnrollment();
   SubmitEnrollmentCredentials();
 
   // Make sure the attribute-prompt view is open.
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepDeviceAttributes);
 
-  WaitForStep("attribute-prompt");
-  EXPECT_TRUE(IsStepDisplayed("attribute-prompt"));
-  EXPECT_FALSE(IsStepDisplayed("success"));
-  EXPECT_FALSE(IsStepDisplayed("error"));
-
-  SubmitAttributePromptUpdate();
-  WaitForStep("success");
+  enrollment_ui_.SubmitDeviceAttributes(test::values::kAssetId,
+                                        test::values::kLocation);
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepSuccess);
 }
 
 // Shows the enrollment screen and mocks the enrollment helper to show license
@@ -517,16 +493,14 @@ IN_PROC_BROWSER_TEST_F(EnterpriseEnrollmentTest, TestLicenseSelection) {
   SubmitEnrollmentCredentials();
 
   // Make sure the license selection screen is open.
-  WaitForStep("license");
-  // Click on third option.
-  test::OobeJS().SelectRadioPath(
-      {"oauth-enroll-license-ui", "license-option-kiosk"});
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepLicenses);
+  // Click on Kiosk option.
+  enrollment_ui_.SelectEnrollmentLicense(test::values::kLicenseTypeKiosk);
   // Click on second option. As there is 0 annual licenses, it should not be
   // selected.
-  test::OobeJS().SelectRadioPath(
-      {"oauth-enroll-license-ui", "license-option-annual"});
-  test::OobeJS().TapOnPath({"oauth-enroll-license-ui", "next"});
-  WaitForStep("success");
+  enrollment_ui_.SelectEnrollmentLicense(test::values::kLicenseTypeAnnual);
+  enrollment_ui_.UseSelectedLicense();
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepSuccess);
 }
 
 // Shows the enrollment screen and mocks the enrollment helper to show Active
@@ -553,12 +527,11 @@ IN_PROC_BROWSER_TEST_F(ActiveDirectoryJoinTest,
   SubmitActiveDirectoryCredentials("machine_name", "" /* machine_dn */, "all",
                                    kAdTestUser, "password");
   WaitForMessage(&message_queue, "\"ShowSpinnerScreen\"");
-  EXPECT_FALSE(IsStepDisplayed("ad-join"));
+  EXPECT_FALSE(enrollment_ui_.IsStepDisplayed(test::ui::kEnrollmentStepADJoin));
 
   CompleteEnrollment();
   // Verify that the success page is displayed.
-  WaitForStep("success");
-  EXPECT_FALSE(IsStepDisplayed("error"));
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepSuccess);
 }
 
 // Verifies that the distinguished name specified on the Active Directory join
@@ -587,12 +560,11 @@ IN_PROC_BROWSER_TEST_F(ActiveDirectoryJoinTest,
                                    "" /* encryption_types */, kAdTestUser,
                                    "password");
   WaitForMessage(&message_queue, "\"ShowSpinnerScreen\"");
-  EXPECT_FALSE(IsStepDisplayed("ad-join"));
+  EXPECT_FALSE(enrollment_ui_.IsStepDisplayed(test::ui::kEnrollmentStepADJoin));
 
   CompleteEnrollment();
   // Verify that the success page is displayed.
-  WaitForStep("success");
-  EXPECT_FALSE(IsStepDisplayed("error"));
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepSuccess);
 }
 
 // Shows the enrollment screen and mocks the enrollment helper to show Active
@@ -616,7 +588,7 @@ IN_PROC_BROWSER_TEST_F(ActiveDirectoryJoinTest,
   SubmitActiveDirectoryCredentials("too_long_machine_name", "" /* machine_dn */,
                                    "" /* encryption_types */, kAdTestUser,
                                    "" /* password */);
-  EXPECT_TRUE(IsStepDisplayed("ad-join"));
+  EXPECT_TRUE(enrollment_ui_.IsStepDisplayed(test::ui::kEnrollmentStepADJoin));
   ExpectElementValid(kAdMachineNameInput, true);
   ExpectElementValid(kAdUsernameInput, true);
   ExpectElementValid(kAdPasswordInput, false);
@@ -626,7 +598,7 @@ IN_PROC_BROWSER_TEST_F(ActiveDirectoryJoinTest,
                                    "" /* encryption_types */, kAdTestUser,
                                    "password");
   WaitForMessage(&message_queue, "\"ShowJoinDomainError\"");
-  EXPECT_TRUE(IsStepDisplayed("ad-join"));
+  EXPECT_TRUE(enrollment_ui_.IsStepDisplayed(test::ui::kEnrollmentStepADJoin));
   ExpectElementValid(kAdMachineNameInput, false);
   ExpectElementValid(kAdUsernameInput, true);
   ExpectElementValid(kAdPasswordInput, true);
@@ -636,7 +608,7 @@ IN_PROC_BROWSER_TEST_F(ActiveDirectoryJoinTest,
                                    "" /* encryption_types */, "test_user",
                                    "password");
   WaitForMessage(&message_queue, "\"ShowJoinDomainError\"");
-  EXPECT_TRUE(IsStepDisplayed("ad-join"));
+  EXPECT_TRUE(enrollment_ui_.IsStepDisplayed(test::ui::kEnrollmentStepADJoin));
   ExpectElementValid(kAdMachineNameInput, true);
   ExpectElementValid(kAdUsernameInput, false);
   ExpectElementValid(kAdPasswordInput, true);
@@ -659,9 +631,9 @@ IN_PROC_BROWSER_TEST_F(ActiveDirectoryJoinTest,
   SubmitActiveDirectoryCredentials("machine_name", "" /* machine_dn */,
                                    "legacy", kAdTestUser, "password");
   WaitForMessage(&message_queue, "\"ShowADJoinError\"");
-  WaitForStep("active-directory-join-error");
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepADJoinError);
   test::OobeJS().TapOnPath({kAdErrorCard, kSubmitButton});
-  WaitForStep("ad-join");
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepADJoin);
 }
 
 // Check that configuration for the streamline Active Directory domain join

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

@@ -31,6 +31,7 @@
 #include "chrome/browser/chromeos/login/existing_user_controller.h"
 #include "chrome/browser/chromeos/login/screens/gaia_view.h"
 #include "chrome/browser/chromeos/login/startup_utils.h"
+#include "chrome/browser/chromeos/login/test/enrollment_ui_mixin.h"
 #include "chrome/browser/chromeos/login/test/fake_gaia_mixin.h"
 #include "chrome/browser/chromeos/login/test/https_forwarder.h"
 #include "chrome/browser/chromeos/login/test/js_checker.h"
@@ -794,11 +795,13 @@ class SAMLEnrollmentTest : public SamlTest {
   guest_view::TestGuestViewManager* GetGuestViewManager();
   content::WebContents* GetEnrollmentContents();
 
- private:
+ protected:
   LocalPolicyTestServerMixin local_policy_mixin_{&mixin_host_};
+  test::EnrollmentUIMixin enrollment_ui_{&mixin_host_};
 
   guest_view::TestGuestViewManagerFactory guest_view_manager_factory_;
 
+ private:
   DISALLOW_COPY_AND_ASSIGN(SAMLEnrollmentTest);
 };
 
@@ -860,7 +863,8 @@ IN_PROC_BROWSER_TEST_F(SAMLEnrollmentTest, WithoutCredentialsPassingAPI) {
   SigninFrameJS().TypeIntoPath("fake_user", {"Email"});
   SigninFrameJS().TypeIntoPath("fake_password", {"Password"});
   SigninFrameJS().TapOn("Submit");
-  OobeBaseTest::WaitForEnrollmentSuccess();
+
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepDeviceAttributes);
 }
 
 IN_PROC_BROWSER_TEST_F(SAMLEnrollmentTest, WithCredentialsPassingAPI) {
@@ -873,7 +877,7 @@ IN_PROC_BROWSER_TEST_F(SAMLEnrollmentTest, WithCredentialsPassingAPI) {
   SigninFrameJS().TypeIntoPath("fake_password", {"Password"});
   SigninFrameJS().TapOn("Submit");
 
-  OobeBaseTest::WaitForEnrollmentSuccess();
+  enrollment_ui_.WaitForStep(test::ui::kEnrollmentStepDeviceAttributes);
 }
 
 class SAMLPolicyTest : public SamlTest {

chrome/browser/chromeos/login/test/enrollment_ui_mixin.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/enrollment_ui_mixin.h"
+
+#include "chrome/browser/chromeos/login/enrollment/enrollment_screen.h"
+#include "chrome/browser/chromeos/login/test/js_checker.h"
+#include "chrome/browser/chromeos/login/test/test_predicate_waiter.h"
+
+namespace chromeos {
+namespace test {
+
+namespace ui {
+
+const char kEnrollmentStepSuccess[] = "success";
+const char kEnrollmentStepError[] = "error";
+const char kEnrollmentStepLicenses[] = "license";
+const char kEnrollmentStepDeviceAttributes[] = "attribute-prompt";
+const char kEnrollmentStepADJoin[] = "ad-join";
+const char kEnrollmentStepADJoinError[] = "active-directory-join-error";
+
+}  // namespace ui
+
+namespace values {
+
+const char kLicenseTypePerpetual[] = "perpetual";
+const char kLicenseTypeAnnual[] = "annual";
+const char kLicenseTypeKiosk[] = "kiosk";
+
+const char kAssetId[] = "asset_id";
+const char kLocation[] = "location";
+
+}  // namespace values
+
+namespace {
+
+const char* const kAllSteps[] = {
+    ui::kEnrollmentStepLicenses, ui::kEnrollmentStepDeviceAttributes,
+    ui::kEnrollmentStepSuccess,  ui::kEnrollmentStepADJoin,
+    ui::kEnrollmentStepError,    ui::kEnrollmentStepADJoinError};
+
+std::string StepVisibleExpression(const std::string& step) {
+  return "document.getElementsByClassName('oauth-enroll-state-" + step +
+         "').length > 0";
+}
+
+}  // namespace
+
+EnrollmentUIMixin::EnrollmentUIMixin(InProcessBrowserTestMixinHost* host)
+    : InProcessBrowserTestMixin(host) {}
+
+EnrollmentUIMixin::~EnrollmentUIMixin() = default;
+
+// Waits until specific enrollment step is displayed.
+void EnrollmentUIMixin::WaitForStep(const std::string& step) {
+  OobeJS().CreateWaiter(StepVisibleExpression(step))->Wait();
+  for (const char* other : kAllSteps) {
+    if (other != step) {
+      ASSERT_FALSE(IsStepDisplayed(other));
+    }
+  }
+}
+// Returns true if there are any DOM elements with the given class.
+bool EnrollmentUIMixin::IsStepDisplayed(const std::string& step) {
+  return OobeJS().GetBool(StepVisibleExpression(step));
+}
+
+void EnrollmentUIMixin::SelectEnrollmentLicense(
+    const std::string& license_type) {
+  OobeJS().SelectRadioPath(
+      {"oauth-enroll-license-ui", "license-option-" + license_type});
+}
+
+void EnrollmentUIMixin::UseSelectedLicense() {
+  OobeJS().TapOnPath({"oauth-enroll-license-ui", "next"});
+}
+
+void EnrollmentUIMixin::SubmitDeviceAttributes(const std::string& asset_id,
+                                               const std::string& location) {
+  OobeJS().TypeIntoPath(asset_id, {"oauth-enroll-asset-id"});
+  OobeJS().TypeIntoPath(location, {"oauth-enroll-location"});
+  OobeJS().TapOn("enroll-attributes-submit-button");
+}
+
+}  // namespace test
+}  // namespace chromeos

chrome/browser/chromeos/login/test/enrollment_ui_mixin.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_ENROLLMENT_UI_MIXIN_H_
+#define CHROME_BROWSER_CHROMEOS_LOGIN_TEST_ENROLLMENT_UI_MIXIN_H_
+
+#include <string>
+
+#include "base/macros.h"
+#include "chrome/browser/chromeos/login/mixin_based_in_process_browser_test.h"
+
+namespace chromeos {
+namespace test {
+
+namespace ui {
+
+//  WaitForStep(...) constants.
+
+extern const char kEnrollmentStepLicenses[];
+extern const char kEnrollmentStepDeviceAttributes[];
+extern const char kEnrollmentStepSuccess[];
+extern const char kEnrollmentStepADJoin[];
+extern const char kEnrollmentStepError[];
+extern const char kEnrollmentStepADJoinError[];
+
+}  // namespace ui
+
+namespace values {
+
+// SelectEnrollmentLicense(...) constants.
+
+extern const char kLicenseTypePerpetual[];
+extern const char kLicenseTypeKiosk[];
+extern const char kLicenseTypeAnnual[];
+
+// SubmitDeviceAttributes common values.
+
+extern const char kAssetId[];
+extern const char kLocation[];
+
+}  // namespace values
+
+// This test mixin covers enrollment-specific OOBE UI interactions.
+class EnrollmentUIMixin : public InProcessBrowserTestMixin {
+ public:
+  explicit EnrollmentUIMixin(InProcessBrowserTestMixinHost* host);
+  ~EnrollmentUIMixin() override;
+
+  // Waits until specific enrollment step is displayed.
+  void WaitForStep(const std::string& step);
+  bool IsStepDisplayed(const std::string& step);
+
+  // Fills out the UI with device attribute information and submits it.
+  void SubmitDeviceAttributes(const std::string& asset_id,
+                              const std::string& location);
+
+  // Selects enrollment license.
+  void SelectEnrollmentLicense(const std::string& license_type);
+
+  // Proceeds with selected license.
+  void UseSelectedLicense();
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(EnrollmentUIMixin);
+};
+
+}  // namespace test
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_CHROMEOS_LOGIN_TEST_ENROLLMENT_UI_MIXIN_H_

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

@@ -58,6 +58,25 @@ void LocalPolicyTestServerMixin::SetUpCommandLine(
                                   policy_test_server_->GetServiceURL().spec());
 }
 
+void LocalPolicyTestServerMixin::ExpectAvailableLicenseCount(int perpetual,
+                                                             int annual,
+                                                             int kiosk) {
+  base::Value licenses(base::Value::Type::DICTIONARY);
+  if (perpetual >= 0) {
+    licenses.SetKey("perpetual", base::Value(perpetual));
+  }
+  if (annual >= 0) {
+    licenses.SetKey("annual", base::Value(annual));
+  }
+  if (kiosk >= 0) {
+    licenses.SetKey("kiosk", base::Value(kiosk));
+  }
+  DCHECK(licenses.DictSize() > 0);
+
+  server_config_.SetKey("available_licenses", std::move(licenses));
+  policy_test_server_->SetConfig(server_config_);
+}
+
 bool LocalPolicyTestServerMixin::UpdateDevicePolicy(
     const enterprise_management::ChromeDeviceSettingsProto& policy) {
   DCHECK(policy_test_server_);

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

@@ -34,6 +34,11 @@ class LocalPolicyTestServerMixin : public InProcessBrowserTestMixin {
   bool UpdateDevicePolicy(
       const enterprise_management::ChromeDeviceSettingsProto& policy);
 
+  // Configures and sets expectations for enrollment flow with license
+  // selection. Non-negative values indicate number of available licenses.
+  // There should be at least one license type.
+  void ExpectAvailableLicenseCount(int perpetual, int annual, int kiosk);
+
   // Set response for DeviceStateRetrievalRequest. Returns that if finds state
   // key passed in the request. State keys could be set by RegisterClient call
   // on policy test server.

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

@@ -173,14 +173,6 @@ void OobeBaseTest::WaitForGaiaPageEvent(const std::string& event) {
   } while (message != "\"Done\"");
 }
 
-void OobeBaseTest::WaitForEnrollmentSuccess() {
-  test::OobeJS()
-      .CreateWaiter(
-          "document.getElementsByClassName('oauth-enroll-state-attribute-"
-          "prompt').length > 0")
-      ->Wait();
-}
-
 void OobeBaseTest::WaitForSigninScreen() {
   WizardController* wizard_controller = WizardController::default_controller();
   if (wizard_controller)

chrome/test/BUILD.gn

@@ -1884,6 +1884,8 @@ test("browser_tests") {
         "../browser/chromeos/login/signin/oauth2_browsertest.cc",
         "../browser/chromeos/login/test/enrollment_helper_mixin.cc",
         "../browser/chromeos/login/test/enrollment_helper_mixin.h",
+        "../browser/chromeos/login/test/enrollment_ui_mixin.cc",
+        "../browser/chromeos/login/test/enrollment_ui_mixin.h",
         "../browser/chromeos/login/test/fake_gaia_mixin.cc",
         "../browser/chromeos/login/test/fake_gaia_mixin.h",
         "../browser/chromeos/login/test/hid_controller_mixin.cc",