Handle Smart Lock for Android promotion screen.

components/proximity_auth/e2e_test/cros.py

Index: components/proximity_auth/e2e_test/cros.py
diff --git a/components/proximity_auth/e2e_test/cros.py b/components/proximity_auth/e2e_test/cros.py
index 4c10226c2cfd94ab24ca2b4cbe5791a12ff5cd9a..1f3c0911448077b283ef42eb31d1912a22755e7b 100644
--- a/components/proximity_auth/e2e_test/cros.py
+++ b/components/proximity_auth/e2e_test/cros.py
@@ -15,11 +15,11 @@ telemetry_dir = os.path.realpath(
 if telemetry_dir not in sys.path:
   sys.path.append(telemetry_dir)
 
-from telemetry.core import browser_options
-from telemetry.core import browser_finder
+from telemetry.internal.browser import browser_options
+from telemetry.internal.browser import browser_finder
 from telemetry.core import exceptions
 from telemetry.core import util
-from telemetry.core.platform import cros_interface
+from telemetry.core import cros_interface
 from telemetry.internal.browser import extension_to_load
 
 logger = logging.getLogger('proximity_auth.%s' % __name__)
@@ -217,6 +217,7 @@ class SmartLockApp(object):
     PAIR = 'pair'
     CLICK_FOR_TRIAL_RUN = 'click_for_trial_run'
     TRIAL_RUN_COMPLETED = 'trial_run_completed'
+    PROMOTE_SMARTLOCK_FOR_ANDROID = 'promote-smart-lock-for-android'
 
   def __init__(self, app_page, chromeos):
     """
@@ -240,6 +241,8 @@ class SmartLockApp(object):
       return SmartLockApp.PairingState.SCAN
     elif state == 'pair':
       return SmartLockApp.PairingState.PAIR
+    elif state == 'promote-smart-lock-for-android':
+      return SmartLockApp.PairingState.PROMOTE_SMARTLOCK_FOR_ANDROID
     elif state == 'complete':
       button_text = self._app_page.EvaluateJavaScript(
           'document.getElementById("pairing-button").textContent')
@@ -282,6 +285,8 @@ class SmartLockApp(object):
     """
     assert(self.pairing_state == self.PairingState.PAIR)
     self._ClickPairingButton()
+    if self.pairing_state == self.PairingState.PROMOTE_SMARTLOCK_FOR_ANDROID:
+      self._ClickPairingButton()
     return self.pairing_state == self.PairingState.CLICK_FOR_TRIAL_RUN
 
   def StartTrialRun(self):
@@ -309,10 +314,14 @@ class SmartLockApp(object):
         'document.getElementById("pairing-button").click()')
 
   def _ClickPairingButton(self):
+    # Waits are needed because the clicks occur before the button label changes.
+    time.sleep(1)
     self._app_page.EvaluateJavaScript(
         'document.getElementById("pairing-button").click()')
+    time.sleep(1)
     util.WaitFor(lambda: self._app_page.EvaluateJavaScript(
         '!document.getElementById("pairing-button").disabled'), 60)
+    time.sleep(1)
     util.WaitFor(lambda: self._app_page.EvaluateJavaScript(
         '!document.getElementById("pairing-button-title")'
         '.classList.contains("animated-fade-out")'), 5)