Added retry policy

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

Index: chrome/browser/chromeos/login/enrollment/enterprise_enrollment_helper_mock.cc
diff --git a/chrome/browser/chromeos/login/enrollment/enterprise_enrollment_helper_mock.cc b/chrome/browser/chromeos/login/enrollment/enterprise_enrollment_helper_mock.cc
index b3a45e2ef0e9c15cc93c1315f838e1b9c34b71e3..de6fd8723fb9597e8cfeac06fd4c1c0fab97e68b 100644
--- a/chrome/browser/chromeos/login/enrollment/enterprise_enrollment_helper_mock.cc
+++ b/chrome/browser/chromeos/login/enrollment/enterprise_enrollment_helper_mock.cc
@@ -2,14 +2,24 @@
 // 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/enrollment/enrollment_screen.h"
 #include "chrome/browser/chromeos/login/enrollment/enterprise_enrollment_helper_mock.h"
 #include "chrome/browser/chromeos/login/enrollment/enterprise_enrollment_helper_impl.h"
+#include "chrome/browser/chromeos/policy/enrollment_status_chromeos.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+using testing::Invoke;
 
 namespace chromeos {
 
 EnterpriseEnrollmentHelperMock::EnterpriseEnrollmentHelperMock(
     EnrollmentStatusConsumer* status_consumer)
-    : EnterpriseEnrollmentHelper(status_consumer) {}
+    : EnterpriseEnrollmentHelper(status_consumer) {
+  ON_CALL(*this, EnrollUsingAttestation())
+      .WillByDefault(
+          Invoke(this, &EnterpriseEnrollmentHelperMock::TryEnrollment));
+}
 
 EnterpriseEnrollmentHelperMock::~EnterpriseEnrollmentHelperMock() {}
 
@@ -18,4 +28,16 @@ EnterpriseEnrollmentHelperMock::status_consumer() {
   return EnterpriseEnrollmentHelper::status_consumer();
 }
 
+void EnterpriseEnrollmentHelperMock::TryEnrollment() {
+  if (should_enroll) {
+    static_cast<EnrollmentScreen*>(status_consumer())
+        ->ShowEnrollmentStatusOnSuccess();
+  } else {
+    status_consumer()->OnEnrollmentError(policy::EnrollmentStatus::ForStatus(
+        policy::EnrollmentStatus::REGISTRATION_FAILED));
+  }
+}
+
+bool EnterpriseEnrollmentHelperMock::should_enroll = true;
+
 }  // namespace chromeos