Clean up affiliation test helper

Changes affiliation test helper from a bunch of static methods to a
class with two factory methods, one for cloud management, the other for
Active Directory management. The former implementation was a bit hand-
waving ("pass the same session manager instance into the two
Set*AffiliationIDs functions", "you don't need fake_auth_policy_client
for non-AD accounts"). The class is much cleaner. Internally, it doesn't
rely on side effects ("we're probably in AD mode if the auth policy
client has been started").

BUG=chromium:865969
TEST=tryjobs

Change-Id: I4ba5e0f9e4b410532cd098cd5ae73f2156a58c83
Reviewed-on: https://chromium-review.googlesource.com/1168496Reviewed-by: Pavol Marko <pmarko@chromium.org>
Reviewed-by: Xiyuan Xia <xiyuan@chromium.org>
Reviewed-by: Devlin <rdevlin.cronin@chromium.org>
Commit-Queue: Lutz Justen <ljusten@chromium.org>
Cr-Commit-Position: refs/heads/master@{#585733}

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

@@ -1078,9 +1078,10 @@ void SAMLPolicyTest::SetUpInProcessBrowserTestFixture() {
   // Initialize device policy.
   std::set<std::string> device_affiliation_ids;
   device_affiliation_ids.insert(kAffiliationID);
-  policy::affiliation_test_helper::SetDeviceAffiliationIDs(
-      &test_helper_, fake_session_manager_client_,
-      nullptr /* fake_auth_policy_client */, device_affiliation_ids);
+  auto affiliation_helper = policy::AffiliationTestHelper::CreateForCloud(
+      fake_session_manager_client_);
+  ASSERT_NO_FATAL_FAILURE((affiliation_helper.SetDeviceAffiliationIDs(
+      &test_helper_, device_affiliation_ids)));
 
   // Initialize user policy.
   EXPECT_CALL(provider_, IsInitializationComplete(_))

chrome/browser/chromeos/policy/affiliation_test_helper.cc

@@ -44,15 +44,12 @@
 
 namespace policy {
 
-namespace affiliation_test_helper {
+namespace {
 
-constexpr char kFakeRefreshToken[] = "fake-refresh-token";
-constexpr char kEnterpriseUserEmail[] = "testuser@example.com";
-constexpr char kEnterpriseUserGaiaId[] = "01234567890";
-
-void SetUserKeys(policy::UserPolicyBuilder* user_policy) {
+// Creates policy key file for the user specified in |user_policy|.
+void SetUserKeys(const policy::UserPolicyBuilder& user_policy) {
   const AccountId account_id =
-      AccountId::FromUserEmail(user_policy->policy_data().username());
+      AccountId::FromUserEmail(user_policy.policy_data().username());
   base::FilePath user_keys_dir;
   ASSERT_TRUE(
       base::PathService::Get(chromeos::DIR_USER_POLICY_KEYS, &user_keys_dir));
@@ -61,7 +58,7 @@ void SetUserKeys(policy::UserPolicyBuilder* user_policy) {
           cryptohome::CreateAccountIdentifierFromAccountId(account_id));
   const base::FilePath user_key_file =
       user_keys_dir.AppendASCII(sanitized_username).AppendASCII("policy.pub");
-  std::string user_key_bits = user_policy->GetPublicSigningKeyAsString();
+  std::string user_key_bits = user_policy.GetPublicSigningKeyAsString();
   ASSERT_FALSE(user_key_bits.empty());
   ASSERT_TRUE(base::CreateDirectory(user_key_file.DirName()));
   ASSERT_EQ(base::WriteFile(user_key_file, user_key_bits.data(),
@@ -69,73 +66,105 @@ void SetUserKeys(policy::UserPolicyBuilder* user_policy) {
             base::checked_cast<int>(user_key_bits.length()));
 }
 
-void SetDeviceAffiliationIDs(
-    policy::DevicePolicyCrosTestHelper* test_helper,
+}  // namespace
+
+constexpr char AffiliationTestHelper::kFakeRefreshToken[] =
+    "fake-refresh-token";
+constexpr char AffiliationTestHelper::kEnterpriseUserEmail[] =
+    "testuser@example.com";
+constexpr char AffiliationTestHelper::kEnterpriseUserGaiaId[] = "01234567890";
+
+// static
+AffiliationTestHelper AffiliationTestHelper::CreateForCloud(
+    chromeos::FakeSessionManagerClient* fake_session_manager_client) {
+  return AffiliationTestHelper(ManagementType::kCloud,
+                               fake_session_manager_client,
+                               nullptr /* fake_auth_policy_client */);
+}
+
+// static
+AffiliationTestHelper AffiliationTestHelper::CreateForActiveDirectory(
+    chromeos::FakeSessionManagerClient* fake_session_manager_client,
+    chromeos::FakeAuthPolicyClient* fake_auth_policy_client) {
+  return AffiliationTestHelper(ManagementType::kActiveDirectory,
+                               fake_session_manager_client,
+                               fake_auth_policy_client);
+}
+
+AffiliationTestHelper::AffiliationTestHelper(AffiliationTestHelper&& other) =
+    default;
+
+AffiliationTestHelper::AffiliationTestHelper(
+    ManagementType management_type,
     chromeos::FakeSessionManagerClient* fake_session_manager_client,
-    chromeos::FakeAuthPolicyClient* fake_auth_policy_client,
+    chromeos::FakeAuthPolicyClient* fake_auth_policy_client)
+    : management_type_(management_type),
+      fake_session_manager_client_(fake_session_manager_client),
+      fake_auth_policy_client_(fake_auth_policy_client) {}
+
+void AffiliationTestHelper::CheckPreconditions() {
+  ASSERT_TRUE(fake_session_manager_client_);
+  ASSERT_TRUE(management_type_ != ManagementType::kActiveDirectory ||
+              fake_auth_policy_client_);
+}
+
+void AffiliationTestHelper::SetDeviceAffiliationIDs(
+    policy::DevicePolicyCrosTestHelper* test_helper,
     const std::set<std::string>& device_affiliation_ids) {
-  // Assume it's Active Directory when the AuthPolicyClient is started.
-  // Make sure we don't overwrite the install attributes in that case.
-  const bool is_active_directory =
-      fake_auth_policy_client && fake_auth_policy_client->started();
-  if (!is_active_directory) {
-    test_helper->InstallOwnerKey();
-    test_helper->MarkAsEnterpriseOwned();
-  }
+  ASSERT_NO_FATAL_FAILURE(CheckPreconditions());
+
   policy::DevicePolicyBuilder* device_policy = test_helper->device_policy();
   for (const auto& device_affiliation_id : device_affiliation_ids) {
     device_policy->policy_data().add_device_affiliation_ids(
         device_affiliation_id);
   }
-  if (!is_active_directory)
+  if (management_type_ != ManagementType::kActiveDirectory) {
+    // Create keys and sign policy. Note that Active Directory policy is
+    // unsigned.
+    test_helper->InstallOwnerKey();
+    test_helper->MarkAsEnterpriseOwned();
     device_policy->SetDefaultSigningKey();
+  }
   device_policy->Build();
 
-  fake_session_manager_client->set_device_policy(device_policy->GetBlob());
-  fake_session_manager_client->OnPropertyChangeComplete(true);
+  fake_session_manager_client_->set_device_policy(device_policy->GetBlob());
+  fake_session_manager_client_->OnPropertyChangeComplete(true);
 
-  // Need fake_auth_policy_client for Active Directory accounts.
-  if (fake_auth_policy_client)
-    fake_auth_policy_client->set_device_affiliation_ids(device_affiliation_ids);
+  if (management_type_ == ManagementType::kActiveDirectory) {
+    fake_auth_policy_client_->set_device_affiliation_ids(
+        device_affiliation_ids);
+  }
 }
 
-void SetUserAffiliationIDs(
+void AffiliationTestHelper::SetUserAffiliationIDs(
     policy::UserPolicyBuilder* user_policy,
-    chromeos::FakeSessionManagerClient* fake_session_manager_client,
-    chromeos::FakeAuthPolicyClient* fake_auth_policy_client,
     const AccountId& user_account_id,
     const std::set<std::string>& user_affiliation_ids) {
-  const bool is_active_directory =
-      user_account_id.GetAccountType() == AccountType::ACTIVE_DIRECTORY;
+  ASSERT_NO_FATAL_FAILURE(CheckPreconditions());
+  ASSERT_TRUE(management_type_ != ManagementType::kActiveDirectory ||
+              user_account_id.GetAccountType() ==
+                  AccountType::ACTIVE_DIRECTORY);
+
   user_policy->policy_data().set_username(user_account_id.GetUserEmail());
-  if (!is_active_directory) {
+  if (management_type_ != ManagementType::kActiveDirectory) {
     user_policy->policy_data().set_gaia_id(user_account_id.GetGaiaId());
-    SetUserKeys(user_policy);
+    ASSERT_NO_FATAL_FAILURE(SetUserKeys(*user_policy));
   }
   for (const auto& user_affiliation_id : user_affiliation_ids) {
     user_policy->policy_data().add_user_affiliation_ids(user_affiliation_id);
   }
   user_policy->Build();
 
-  // Make sure AD policy is stored using the proper cryptohome key. This code
-  // runs before the AD account is migrated, so it would use the email address.
-  // TODO(ljusten): Clean this up as soon as CL:1055509 lands.
-  cryptohome::Identification cryptohome_id =
-      is_active_directory ? cryptohome::Identification::FromString(
-                                user_account_id.GetAccountIdKey())
-                          : cryptohome::Identification(user_account_id);
-
-  fake_session_manager_client->set_user_policy(
-      cryptohome::CreateAccountIdentifierFromIdentification(cryptohome_id),
+  fake_session_manager_client_->set_user_policy(
+      cryptohome::CreateAccountIdentifierFromAccountId(user_account_id),
       user_policy->GetBlob());
 
-  // Need fake_auth_policy_client for Active Directory accounts.
-  DCHECK(!is_active_directory || fake_auth_policy_client);
-  if (fake_auth_policy_client)
-    fake_auth_policy_client->set_user_affiliation_ids(user_affiliation_ids);
+  if (management_type_ == ManagementType::kActiveDirectory)
+    fake_auth_policy_client_->set_user_affiliation_ids(user_affiliation_ids);
 }
 
-void PreLoginUser(const AccountId& account_id) {
+// static
+void AffiliationTestHelper::PreLoginUser(const AccountId& account_id) {
   ListPrefUpdate users_pref(g_browser_process->local_state(), "LoggedInUsers");
   users_pref->AppendIfNotPresent(
       std::make_unique<base::Value>(account_id.GetUserEmail()));
@@ -145,7 +174,8 @@ void PreLoginUser(const AccountId& account_id) {
   chromeos::StartupUtils::MarkOobeCompleted();
 }
 
-void LoginUser(const AccountId& account_id) {
+// static
+void AffiliationTestHelper::LoginUser(const AccountId& account_id) {
   chromeos::test::UserSessionManagerTestApi session_manager_test_api(
       chromeos::UserSessionManager::GetInstance());
   session_manager_test_api.SetShouldObtainTokenHandleInTests(false);
@@ -180,7 +210,9 @@ void LoginUser(const AccountId& account_id) {
                 << " was not added via PreLoginUser()";
 }
 
-void AppendCommandLineSwitchesForLoginManager(base::CommandLine* command_line) {
+// static
+void AffiliationTestHelper::AppendCommandLineSwitchesForLoginManager(
+    base::CommandLine* command_line) {
   command_line->AppendSwitch(chromeos::switches::kLoginManager);
   command_line->AppendSwitch(chromeos::switches::kForceLoginManagerInTests);
   // LoginManager tests typically don't stand up a policy test server but
@@ -190,6 +222,4 @@ void AppendCommandLineSwitchesForLoginManager(base::CommandLine* command_line) {
       chromeos::switches::kAllowFailedPolicyFetchForTest);
 }
 
-}  // namespace affiliation_test_helper
-
 }  // namespace policy

chrome/browser/chromeos/policy/affiliation_test_helper.h

@@ -7,6 +7,8 @@
 
 #include <set>
 #include <string>
+
+#include "base/macros.h"
 #include "components/policy/core/common/cloud/policy_builder.h"
 
 class AccountId;
@@ -24,102 +26,77 @@ namespace policy {
 
 class DevicePolicyCrosTestHelper;
 
-namespace affiliation_test_helper {
-
-// Creates policy key file for the user specified in |user_policy|.
-// TODO(peletskyi): Replace pointer with const reference and replace this
-// boilerplate in other places (http://crbug.com/549111).
-void SetUserKeys(policy::UserPolicyBuilder* user_policy);
-
-// Sets device affiliation ID to |fake_session_manager_client| from
-// |device_affiliation_ids| and modifies |test_helper| so that it contains
-// correct values of device affiliation IDs for future use. To add some device
-// policies and have device affiliation ID valid please use |test_helper|
-// modified by this function. Example:
-//
-// FakeSessionManagerClient* fake_session_manager_client =
-//   new chromeos::FakeSessionManagerClient;
-// DBusThreadManager::GetSetterForTesting()->SetSessionManagerClient(
-//   std::unique_ptr<SessionManagerClient>(fake_session_manager_client));
-//
-// policy::DevicePolicyCrosTestHelper test_helper;
-// std::set<std::string> device_affiliation_ids;
-// device_affiliation_ids.insert(some-affiliation-id);
-//
-// affiliation_test_helper::SetDeviceAffiliationIDs(
-//     &test_helper,
-//     fake_session_manager_client,
-//     nullptr /* fake_auth_policy_client */
-//     device_affiliation_ids);
-//
-// If it is used together with SetUserAffiliationIDs() (which is the most common
-// case) |fake_session_manager_client| must point to the same object as in
-// SetUserAffiliationIDs() call.
-// In browser tests one can call this function from
-// SetUpInProcessBrowserTestFixture().
-//
-// |fake_auth_policy_client| is only needed if the test supports Active
-// Directory user accounts. If not, it may be nullptr.
-void SetDeviceAffiliationIDs(
-    policy::DevicePolicyCrosTestHelper* test_helper,
-    chromeos::FakeSessionManagerClient* fake_session_manager_client,
-    chromeos::FakeAuthPolicyClient* fake_auth_policy_client,
-    const std::set<std::string>& device_affiliation_ids);
-
-// Sets user affiliation ID for |user_account_id| to
-// |fake_session_manager_client| from |user_affiliation_ids| and modifies
-// |user_policy| so that it contains correct values of user affiliation IDs for
-// future use. To add user policies and have user affiliation IDs valid please
-// use |user_policy| modified by this function. Example:
-//
-// FakeSessionManagerClient* fake_session_manager_client =
-//    new chromeos::FakeSessionManagerClient;
-// DBusThreadManager::GetSetterForTesting()->SetSessionManagerClient(
-//    std::unique_ptr<SessionManagerClient>(fake_session_manager_client));
-//
-// policy::UserPolicyBuilder user_policy;
-// std::set<std::string> user_affiliation_ids;
-// user_affiliation_ids.insert("some-affiliation-id");
-//
-// affiliation_test_helper::SetUserAffiliationIDs(
-//    &user_policy,
-//    fake_session_manager_client,
-//    nullptr /* fake_auth_policy_client */,
-//    account_id,
-//    user_affiliation_ids);
-//
-// If it is used together SetDeviceAffiliationIDs() (which is the most common
-// case) |fake_session_manager_client| must point to the same object as in
-// SetDeviceAffiliationIDs() call.
-// In browser tests one can call this function from
-// SetUpInProcessBrowserTestFixture().
-//
-// |fake_auth_policy_client| is only needed if the test supports Active
-// Directory user accounts. If not, it may be nullptr.
-void SetUserAffiliationIDs(
-    policy::UserPolicyBuilder* user_policy,
-    chromeos::FakeSessionManagerClient* fake_session_manager_client,
-    chromeos::FakeAuthPolicyClient* fake_auth_policy_client,
-    const AccountId& user_account_id,
-    const std::set<std::string>& user_affiliation_ids);
-
-// Registers the user with the given |account_id| on the device and marks OOBE
-// as completed. This method should be called in PRE_* test.
-void PreLoginUser(const AccountId& account_id);
-
-// Log in user with |account_id|. User should be registered using
-// PreLoginUser().
-void LoginUser(const AccountId& user_id);
-
-// Set necessary for login command line switches. Execute it in
-// SetUpCommandLine().
-void AppendCommandLineSwitchesForLoginManager(base::CommandLine* command_line);
-
-extern const char kFakeRefreshToken[];
-extern const char kEnterpriseUserEmail[];
-extern const char kEnterpriseUserGaiaId[];
-
-}  // namespace affiliation_test_helper
+class AffiliationTestHelper {
+ public:
+  // Creates an |AffiliationTestHelper| for Cloud management (regular Google
+  // accounts). The |fake_session_manager_client| pointer must outlive this
+  // object.
+  static AffiliationTestHelper CreateForCloud(
+      chromeos::FakeSessionManagerClient* fake_session_manager_client);
+
+  // Creates an |AffiliationTestHelper| for Active Directory management (Active
+  // Directory accounts). The pointers must outlive this object.
+  static AffiliationTestHelper CreateForActiveDirectory(
+      chromeos::FakeSessionManagerClient* fake_session_manager_client,
+      chromeos::FakeAuthPolicyClient* fake_auth_policy_client);
+
+  // Allow move construction, so the static constructors can be used despite
+  // DISALLOW_COPY_AND_ASSIGN.
+  AffiliationTestHelper(AffiliationTestHelper&& other);
+
+  // Sets device affiliation IDs to |device_affiliation_ids| in
+  // |fake_session_manager_client_| and modifies |test_helper| so that it
+  // contains correct values of device affiliation IDs for future use. To add
+  // some device policies and have device affiliation ID valid use |test_helper|
+  // modified by this function.
+  void SetDeviceAffiliationIDs(
+      policy::DevicePolicyCrosTestHelper* test_helper,
+      const std::set<std::string>& device_affiliation_ids);
+
+  // Sets user affiliation IDs to |user_affiliation_ids| in
+  // |fake_session_manager_client| and modifies |user_policy| so that it
+  // contains correct values of user affiliation IDs for future use. To add user
+  // policies and have user affiliation IDs valid please use |user_policy|
+  // modified by this function.
+  void SetUserAffiliationIDs(policy::UserPolicyBuilder* user_policy,
+                             const AccountId& user_account_id,
+                             const std::set<std::string>& user_affiliation_ids);
+
+  // Registers the user with the given |account_id| on the device and marks OOBE
+  // as completed. This method should be called in PRE_* test.
+  static void PreLoginUser(const AccountId& account_id);
+
+  // Log in user with |account_id|. User should be registered using
+  // PreLoginUser().
+  static void LoginUser(const AccountId& user_id);
+
+  // Set necessary for login command line switches. Execute it in
+  // SetUpCommandLine().
+  static void AppendCommandLineSwitchesForLoginManager(
+      base::CommandLine* command_line);
+
+  static const char kFakeRefreshToken[];
+  static const char kEnterpriseUserEmail[];
+  static const char kEnterpriseUserGaiaId[];
+
+ private:
+  enum class ManagementType { kCloud, kActiveDirectory };
+
+  AffiliationTestHelper(
+      ManagementType management_type,
+      chromeos::FakeSessionManagerClient* fake_session_manager_client,
+      chromeos::FakeAuthPolicyClient* fake_auth_policy_client);
+
+  // ASSERTs on pointer validity.
+  void CheckPreconditions();
+
+  ManagementType management_type_;
+  chromeos::FakeSessionManagerClient*
+      fake_session_manager_client_;                          // Not owned.
+  chromeos::FakeAuthPolicyClient* fake_auth_policy_client_;  // Not owned.
+
+  DISALLOW_COPY_AND_ASSIGN(AffiliationTestHelper);
+};
 
 }  // namespace policy
 

chrome/browser/chromeos/policy/unaffiliated_arc_allowed_browsertest.cc

@@ -56,7 +56,7 @@ class UnaffiliatedArcAllowedTest
   void SetUpCommandLine(base::CommandLine* command_line) override {
     DevicePolicyCrosBrowserTest::SetUpCommandLine(command_line);
     arc::SetArcAvailableCommandLineForTesting(command_line);
-    affiliation_test_helper::AppendCommandLineSwitchesForLoginManager(
+    AffiliationTestHelper::AppendCommandLineSwitchesForLoginManager(
         command_line);
   }
 
@@ -70,14 +70,12 @@ class UnaffiliatedArcAllowedTest
     const std::set<std::string> user_affiliation_ids = {
         GetParam().affiliated ? kAffiliationID : kAnotherAffiliationID};
 
-    affiliation_test_helper::SetDeviceAffiliationIDs(
-        &test_helper, session_manager_client(),
-        nullptr /* fake_auth_policy_client */, device_affiliation_ids);
-
-    affiliation_test_helper::SetUserAffiliationIDs(
-        &user_policy, session_manager_client(),
-        nullptr /* fake_auth_policy_client */, affiliated_account_id_,
-        user_affiliation_ids);
+    AffiliationTestHelper affiliation_helper =
+        AffiliationTestHelper::CreateForCloud(session_manager_client());
+    ASSERT_NO_FATAL_FAILURE(affiliation_helper.SetDeviceAffiliationIDs(
+        &test_helper, device_affiliation_ids));
+    ASSERT_NO_FATAL_FAILURE(affiliation_helper.SetUserAffiliationIDs(
+        &user_policy, affiliated_account_id_, user_affiliation_ids));
   }
 
   void TearDownOnMainThread() override {
@@ -115,11 +113,11 @@ class UnaffiliatedArcAllowedTest
 };
 
 IN_PROC_BROWSER_TEST_P(UnaffiliatedArcAllowedTest, PRE_ProfileTest) {
-  affiliation_test_helper::PreLoginUser(affiliated_account_id_);
+  AffiliationTestHelper::PreLoginUser(affiliated_account_id_);
 }
 
 IN_PROC_BROWSER_TEST_P(UnaffiliatedArcAllowedTest, ProfileTest) {
-  affiliation_test_helper::LoginUser(affiliated_account_id_);
+  AffiliationTestHelper::LoginUser(affiliated_account_id_);
   const user_manager::User* user =
       user_manager::UserManager::Get()->FindUser(affiliated_account_id_);
   const Profile* profile =

chrome/browser/chromeos/policy/user_affiliation_browsertest.cc

@@ -136,7 +136,7 @@ class UserAffiliationBrowserTest
   void SetUpCommandLine(base::CommandLine* command_line) override {
     InProcessBrowserTest::SetUpCommandLine(command_line);
     if (content::IsPreTest()) {
-      affiliation_test_helper::AppendCommandLineSwitchesForLoginManager(
+      AffiliationTestHelper::AppendCommandLineSwitchesForLoginManager(
           command_line);
     } else {
       const cryptohome::AccountIdentifier cryptohome_id =
@@ -182,13 +182,18 @@ class UserAffiliationBrowserTest
     const std::set<std::string> user_affiliation_ids = {
         GetParam().affiliated ? kAffiliationID : kAnotherAffiliationID};
 
-    affiliation_test_helper::SetDeviceAffiliationIDs(
-        &test_helper, fake_session_manager_client, fake_auth_policy_client,
-        device_affiliation_ids);
+    AffiliationTestHelper affiliation_helper =
+        GetParam().active_directory
+            ? AffiliationTestHelper::CreateForActiveDirectory(
+                  fake_session_manager_client, fake_auth_policy_client)
+            : AffiliationTestHelper::CreateForCloud(
+                  fake_session_manager_client);
 
-    affiliation_test_helper::SetUserAffiliationIDs(
-        &user_policy, fake_session_manager_client, fake_auth_policy_client,
-        account_id_, user_affiliation_ids);
+    ASSERT_NO_FATAL_FAILURE(affiliation_helper.SetDeviceAffiliationIDs(
+        &test_helper, device_affiliation_ids));
+
+    ASSERT_NO_FATAL_FAILURE(affiliation_helper.SetUserAffiliationIDs(
+        &user_policy, account_id_, user_affiliation_ids));
 
     // Set retry delay to prevent timeouts.
     policy::DeviceManagementService::SetRetryDelayForTesting(0);
@@ -274,12 +279,12 @@ class UserAffiliationBrowserTest
 };
 
 IN_PROC_BROWSER_TEST_P(UserAffiliationBrowserTest, PRE_PRE_TestAffiliation) {
-  affiliation_test_helper::PreLoginUser(account_id_);
+  AffiliationTestHelper::PreLoginUser(account_id_);
 }
 
 // This part of the test performs a regular sign-in through the login manager.
 IN_PROC_BROWSER_TEST_P(UserAffiliationBrowserTest, PRE_TestAffiliation) {
-  affiliation_test_helper::LoginUser(account_id_);
+  AffiliationTestHelper::LoginUser(account_id_);
   ASSERT_NO_FATAL_FAILURE(VerifyAffiliationExpectations());
 }
 

chrome/browser/extensions/api/enterprise_device_attributes/enterprise_device_attributes_apitest.cc

@@ -100,7 +100,7 @@ class EnterpriseDeviceAttributesTest
   // ExtensionApiTest
   void SetUpCommandLine(base::CommandLine* command_line) override {
     ExtensionApiTest::SetUpCommandLine(command_line);
-    policy::affiliation_test_helper::AppendCommandLineSwitchesForLoginManager(
+    policy::AffiliationTestHelper::AppendCommandLineSwitchesForLoginManager(
         command_line);
   }
 
@@ -113,11 +113,14 @@ class EnterpriseDeviceAttributesTest
         std::unique_ptr<chromeos::SessionManagerClient>(
             fake_session_manager_client));
 
+    policy::AffiliationTestHelper affiliation_helper =
+        policy::AffiliationTestHelper::CreateForCloud(
+            fake_session_manager_client);
+
     std::set<std::string> device_affiliation_ids;
     device_affiliation_ids.insert(kAffiliationID);
-    policy::affiliation_test_helper::SetDeviceAffiliationIDs(
-        &test_helper_, fake_session_manager_client,
-        nullptr /* fake_auth_policy_client */, device_affiliation_ids);
+    ASSERT_NO_FATAL_FAILURE(affiliation_helper.SetDeviceAffiliationIDs(
+        &test_helper_, device_affiliation_ids));
 
     std::set<std::string> user_affiliation_ids;
     if (GetParam().affiliated) {
@@ -126,10 +129,8 @@ class EnterpriseDeviceAttributesTest
       user_affiliation_ids.insert(kAnotherAffiliationID);
     }
     policy::UserPolicyBuilder user_policy;
-    policy::affiliation_test_helper::SetUserAffiliationIDs(
-        &user_policy, fake_session_manager_client,
-        nullptr /* fake_auth_policy_client */, affiliated_account_id_,
-        user_affiliation_ids);
+    ASSERT_NO_FATAL_FAILURE(affiliation_helper.SetUserAffiliationIDs(
+        &user_policy, affiliated_account_id_, user_affiliation_ids));
 
     test_helper_.InstallOwnerKey();
     // Init the device policy.
@@ -158,7 +159,7 @@ class EnterpriseDeviceAttributesTest
     const base::ListValue* users =
         g_browser_process->local_state()->GetList("LoggedInUsers");
     if (!users->empty())
-      policy::affiliation_test_helper::LoginUser(affiliated_account_id_);
+      policy::AffiliationTestHelper::LoginUser(affiliated_account_id_);
 
     ExtensionApiTest::SetUpOnMainThread();
   }
@@ -231,7 +232,7 @@ class EnterpriseDeviceAttributesTest
 };
 
 IN_PROC_BROWSER_TEST_P(EnterpriseDeviceAttributesTest, PRE_Success) {
-  policy::affiliation_test_helper::PreLoginUser(affiliated_account_id_);
+  policy::AffiliationTestHelper::PreLoginUser(affiliated_account_id_);
 }
 
 IN_PROC_BROWSER_TEST_P(EnterpriseDeviceAttributesTest, Success) {

chrome/browser/extensions/api/platform_keys/platform_keys_test_base.cc

@@ -34,9 +34,6 @@
 const char kAffiliationID[] = "some-affiliation-id";
 const char kTestUserinfoToken[] = "fake-userinfo-token";
 
-using policy::affiliation_test_helper::kEnterpriseUserEmail;
-using policy::affiliation_test_helper::kEnterpriseUserGaiaId;
-
 PlatformKeysTestBase::PlatformKeysTestBase(
     SystemTokenStatus system_token_status,
     EnrollmentStatus enrollment_status,
@@ -44,8 +41,9 @@ PlatformKeysTestBase::PlatformKeysTestBase(
     : system_token_status_(system_token_status),
       enrollment_status_(enrollment_status),
       user_status_(user_status),
-      account_id_(AccountId::FromUserEmailGaiaId(kEnterpriseUserEmail,
-                                                 kEnterpriseUserGaiaId)) {
+      account_id_(AccountId::FromUserEmailGaiaId(
+          policy::AffiliationTestHelper::kEnterpriseUserEmail,
+          policy::AffiliationTestHelper::kEnterpriseUserGaiaId)) {
   // Command line should not be tweaked as if user is already logged in.
   set_chromeos_user_ = false;
   // We log in without running browser.
@@ -78,7 +76,7 @@ void PlatformKeysTestBase::SetUp() {
 void PlatformKeysTestBase::SetUpCommandLine(base::CommandLine* command_line) {
   extensions::ExtensionApiTest::SetUpCommandLine(command_line);
 
-  policy::affiliation_test_helper::AppendCommandLineSwitchesForLoginManager(
+  policy::AffiliationTestHelper::AppendCommandLineSwitchesForLoginManager(
       command_line);
 
   const GURL gaia_url = gaia_https_forwarder_.GetURLForSSLHost(std::string());
@@ -99,22 +97,23 @@ void PlatformKeysTestBase::SetUpInProcessBrowserTestFixture() {
       std::unique_ptr<chromeos::SessionManagerClient>(
           fake_session_manager_client));
 
+  policy::AffiliationTestHelper affiliation_helper =
+      policy::AffiliationTestHelper::CreateForCloud(
+          fake_session_manager_client);
+
   if (enrollment_status() == EnrollmentStatus::ENROLLED) {
     std::set<std::string> device_affiliation_ids;
     device_affiliation_ids.insert(kAffiliationID);
-    policy::affiliation_test_helper::SetDeviceAffiliationIDs(
-        &device_policy_test_helper_, fake_session_manager_client,
-        nullptr /* fake_auth_policy_client */, device_affiliation_ids);
+    ASSERT_NO_FATAL_FAILURE(affiliation_helper.SetDeviceAffiliationIDs(
+        &device_policy_test_helper_, device_affiliation_ids));
   }
 
   if (user_status() == UserStatus::MANAGED_AFFILIATED_DOMAIN) {
     std::set<std::string> user_affiliation_ids;
     user_affiliation_ids.insert(kAffiliationID);
     policy::UserPolicyBuilder user_policy;
-    policy::affiliation_test_helper::SetUserAffiliationIDs(
-        &user_policy, fake_session_manager_client,
-        nullptr /* fake_auth_policy_client */, account_id_,
-        user_affiliation_ids);
+    ASSERT_NO_FATAL_FAILURE(affiliation_helper.SetUserAffiliationIDs(
+        &user_policy, account_id_, user_affiliation_ids));
   }
 
   EXPECT_CALL(mock_policy_provider_, IsInitializationComplete(testing::_))
@@ -136,7 +135,7 @@ void PlatformKeysTestBase::SetUpOnMainThread() {
   token_info.audience = GaiaUrls::GetInstance()->oauth2_chrome_client_id();
   token_info.token = kTestUserinfoToken;
   token_info.email = account_id_.GetUserEmail();
-  fake_gaia_.IssueOAuthToken(policy::affiliation_test_helper::kFakeRefreshToken,
+  fake_gaia_.IssueOAuthToken(policy::AffiliationTestHelper::kFakeRefreshToken,
                              token_info);
 
   // On PRE_ test stage list of users is empty at this point. Then in the body
@@ -144,7 +143,7 @@ void PlatformKeysTestBase::SetUpOnMainThread() {
   // after PRE_ test the list of user contains one kEnterpriseUser user.
   // This user logs in.
   if (!IsPreTest()) {
-    policy::affiliation_test_helper::LoginUser(account_id_);
+    policy::AffiliationTestHelper::LoginUser(account_id_);
 
     if (user_status() != UserStatus::UNMANAGED) {
       policy::ProfilePolicyConnector* const connector =
@@ -184,7 +183,7 @@ void PlatformKeysTestBase::PrepareTestSystemSlotOnIO(
     crypto::ScopedTestSystemNSSKeySlot* system_slot) {}
 
 void PlatformKeysTestBase::RunPreTest() {
-  policy::affiliation_test_helper::PreLoginUser(account_id_);
+  policy::AffiliationTestHelper::PreLoginUser(account_id_);
 }
 
 bool PlatformKeysTestBase::TestExtension(const std::string& page_url) {