Switch to id-based affiliation determination

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

Index: chrome/browser/chromeos/login/saml/saml_browsertest.cc
diff --git a/chrome/browser/chromeos/login/saml/saml_browsertest.cc b/chrome/browser/chromeos/login/saml/saml_browsertest.cc
index 25f1e7dfe3019534bc390783092dd3f0ab8b4028..fb3ca9ce47c1fc13bb8ea4e9068871b4394fab17 100644
--- a/chrome/browser/chromeos/login/saml/saml_browsertest.cc
+++ b/chrome/browser/chromeos/login/saml/saml_browsertest.cc
@@ -31,7 +31,9 @@
 #include "chrome/browser/chromeos/login/test/oobe_screen_waiter.h"
 #include "chrome/browser/chromeos/login/ui/login_display_host_impl.h"
 #include "chrome/browser/chromeos/login/ui/webui_login_display.h"
+#include "chrome/browser/chromeos/login/users/chrome_user_manager.h"
 #include "chrome/browser/chromeos/login/wizard_controller.h"
+#include "chrome/browser/chromeos/policy/affiliation_test_helper.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/policy/proto/chrome_device_policy.pb.h"
@@ -133,6 +135,8 @@ const char kTestUserinfoToken[] = "fake-userinfo-token";
 const char kTestRefreshToken[] = "fake-refresh-token";
 const char kPolicy[] = "{\"managed_users\": [\"*\"]}";
 
+const char kAffiliationID[] = "some-affiliation-id";
+
 // FakeSamlIdp serves IdP auth form and the form submission. The form is
 // served with the template's RelayState placeholder expanded to the real
 // RelayState parameter from request. The form submission redirects back to
@@ -951,12 +955,10 @@ void SAMLPolicyTest::SetUpInProcessBrowserTestFixture() {
   SamlTest::SetUpInProcessBrowserTestFixture();
 
   // Initialize device policy.
-  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);
+  std::set<std::string> device_affiliation_ids;
+  device_affiliation_ids.insert(kAffiliationID);
+  policy::affiliation_test_helper::SetDeviceAffiliationID(
+      &test_helper_, fake_session_manager_client_, device_affiliation_ids);
 
   // Initialize user policy.
   EXPECT_CALL(provider_, IsInitializationComplete(_))
@@ -978,6 +980,18 @@ void SAMLPolicyTest::SetUpOnMainThread() {
       AccountId::FromUserEmail(kDifferentDomainSAMLUserEmail),
       user_manager::User::OAUTH2_TOKEN_STATUS_VALID);
 
+  // Give affiliated users appropriate affiliation IDs.
+  std::set<std::string> user_affiliation_ids;
+  user_affiliation_ids.insert(kAffiliationID);
+  chromeos::ChromeUserManager::Get()->SetUserAffiliation(kFirstSAMLUserEmail,
+                                                         user_affiliation_ids);
+  chromeos::ChromeUserManager::Get()->SetUserAffiliation(kSecondSAMLUserEmail,
+                                                         user_affiliation_ids);
+  chromeos::ChromeUserManager::Get()->SetUserAffiliation(kHTTPSAMLUserEmail,
+                                                         user_affiliation_ids);
+  chromeos::ChromeUserManager::Get()->SetUserAffiliation(kNonSAMLUserEmail,
+                                                         user_affiliation_ids);
+
   // Set up fake networks.
   DBusThreadManager::Get()
       ->GetShillManagerClient()