Enhance bootloader recovery handling.

Poll and retry for device to enter known state.

Bug 6259681

Change-Id: Ifba3264fa1c81a4afc76b2cc2d407ac7f6b8f043
diff --git a/src/com/android/tradefed/device/WaitDeviceRecovery.java b/src/com/android/tradefed/device/WaitDeviceRecovery.java
index cf69b82..0826cbe 100644
--- a/src/com/android/tradefed/device/WaitDeviceRecovery.java
+++ b/src/com/android/tradefed/device/WaitDeviceRecovery.java
@@ -37,6 +37,8 @@
     /** the time in ms to wait before beginning recovery attempts */
     protected static final long INITIAL_PAUSE_TIME = 5 * 1000;
 
+    private static final long BOOTLOADER_POLL_ATTEMPTS = 3;
+
     // TODO: add a separate configurable timeout per operation
     @Option(name="device-wait-time",
             description="maximum time in ms to wait for a single device recovery command.")
@@ -159,28 +161,42 @@
                 INITIAL_PAUSE_TIME, monitor.getSerialNumber()));
         getRunUtil().sleep(INITIAL_PAUSE_TIME);
 
-        // ensure bootloader state is updated
-        monitor.waitForDeviceBootloaderStateUpdate();
+        // poll and wait for device to return to valid state
+        long pollTime = mBootloaderWaitTime / BOOTLOADER_POLL_ATTEMPTS;
+        for (int i=0; i < BOOTLOADER_POLL_ATTEMPTS; i++) {
+            if (monitor.waitForDeviceBootloader(pollTime)) {
+                handleDeviceBootloaderUnresponsive(monitor);
+                // passed above check, abort
+                return;
+            } else if (monitor.getDeviceState() == TestDeviceState.ONLINE) {
+                handleDeviceOnlineExpectedBootloader(monitor);
+                return;
+            }
+        }
+        handleDeviceBootloaderNotAvailable(monitor);
+    }
 
-        if (monitor.getDeviceState().equals(TestDeviceState.FASTBOOT)) {
-            CLog.i("Found device %s in fastboot but potentially unresponsive.",
-                    monitor.getSerialNumber());
-            handleDeviceBootloaderUnresponsive(monitor);
-        } else {
-            if (monitor.getDeviceState() == TestDeviceState.ONLINE) {
-                Log.i(LOG_TAG, String.format(
-                    "Found device %s online but expected fastboot.",
-                    monitor.getSerialNumber()));
-                IDevice device = monitor.waitForDeviceOnline();
-                if (device == null) {
-                    handleDeviceBootloaderNotAvailable(monitor);
-                    return;
-                }
-                rebootDeviceIntoBootloader(device);
-            }
-            if (!monitor.waitForDeviceBootloader(mBootloaderWaitTime)) {
-                handleDeviceBootloaderNotAvailable(monitor);
-            }
+    /**
+     * Handle condition where device is online, but should be in bootloader state.
+     * <p/>
+     * If this method
+     * @param monitor
+     * @throws DeviceNotAvailableException
+     */
+    protected void handleDeviceOnlineExpectedBootloader(final IDeviceStateMonitor monitor)
+            throws DeviceNotAvailableException {
+        Log.i(LOG_TAG, String.format("Found device %s online but expected fastboot.",
+            monitor.getSerialNumber()));
+        // call waitForDeviceOnline to get handle to IDevice
+        IDevice device = monitor.waitForDeviceOnline();
+        if (device == null) {
+            handleDeviceBootloaderNotAvailable(monitor);
+            return;
+        }
+        rebootDeviceIntoBootloader(device);
+        if (!monitor.waitForDeviceBootloader(mBootloaderWaitTime)) {
+            throw new DeviceNotAvailableException(String.format(
+                    "Device %s not in bootloader after reboot", monitor.getSerialNumber()));
         }
     }
 
@@ -190,12 +206,14 @@
      */
     protected void handleDeviceBootloaderUnresponsive(IDeviceStateMonitor monitor)
             throws DeviceNotAvailableException {
+        CLog.i("Found device %s in fastboot but potentially unresponsive.",
+                monitor.getSerialNumber());
         // TODO: retry reboot
         getRunUtil().runTimedCmd(20*1000, "fastboot", "-s", monitor.getSerialNumber(),
                 "reboot-bootloader");
         // wait for device to reboot
         monitor.waitForDeviceNotAvailable(20*1000);
-        if (!monitor.waitForDeviceBootloader(mWaitTime)) {
+        if (!monitor.waitForDeviceBootloader(mBootloaderWaitTime)) {
             throw new DeviceNotAvailableException(String.format(
                     "Device %s not in bootloader after reboot", monitor.getSerialNumber()));
         }
diff --git a/tests/res/config/tf/func.xml b/tests/res/config/tf/func.xml
index 7aeee45..aa10a1b 100644
--- a/tests/res/config/tf/func.xml
+++ b/tests/res/config/tf/func.xml
@@ -20,7 +20,7 @@
         <option name="apk-name" value="TradeFedTestApp.apk" />
         <option name="apk-name" value="TradeFedUiTestApp.apk" />
     </target_preparer>
-    <test class="com.android.tradefed.FuncTests" />
+    <test class="com.android.tradefed.testtype.HostTest" />
     <logger class="com.android.tradefed.log.FileLogger" >
         <option name="log-level-display" value="debug" />
     </logger>
diff --git a/tests/src/com/android/tradefed/device/WaitDeviceRecoveryTest.java b/tests/src/com/android/tradefed/device/WaitDeviceRecoveryTest.java
index 74e8ffc..c4fa088 100644
--- a/tests/src/com/android/tradefed/device/WaitDeviceRecoveryTest.java
+++ b/tests/src/com/android/tradefed/device/WaitDeviceRecoveryTest.java
@@ -144,18 +144,14 @@
      */
     public void testRecoverDeviceBootloader_fastboot() throws DeviceNotAvailableException {
         mMockRunUtil.sleep(EasyMock.anyLong());
-        mMockMonitor.waitForDeviceBootloaderStateUpdate();
-        EasyMock.expect(mMockMonitor.getDeviceState()).andReturn(TestDeviceState.FASTBOOT);
-        CommandResult result = new CommandResult();
-        result.setStatus(CommandStatus.SUCCESS);
         // expect reboot
         EasyMock.expect(mMockRunUtil.runTimedCmd(EasyMock.anyLong(), EasyMock.eq("fastboot"),
                 EasyMock.eq("-s"), EasyMock.eq("serial"), EasyMock.eq("reboot-bootloader"))).
-                andReturn(result);
+                andReturn(new CommandResult(CommandStatus.SUCCESS));
         EasyMock.expect(mMockMonitor.waitForDeviceNotAvailable(EasyMock.anyLong())).andReturn(
                 Boolean.TRUE);
         EasyMock.expect(mMockMonitor.waitForDeviceBootloader(EasyMock.anyLong())).andReturn(
-                Boolean.TRUE);
+                Boolean.TRUE).times(2);
         replayMocks();
         mRecovery.recoverDeviceBootloader(mMockMonitor);
         verifyMocks();
@@ -167,11 +163,17 @@
      */
     public void testRecoverDeviceBootloader_unavailable() throws DeviceNotAvailableException {
         mMockRunUtil.sleep(EasyMock.anyLong());
-        mMockMonitor.waitForDeviceBootloaderStateUpdate();
-        EasyMock.expect(mMockMonitor.getDeviceState()).andReturn(TestDeviceState.NOT_AVAILABLE).
-                times(2);
         EasyMock.expect(mMockMonitor.waitForDeviceBootloader(EasyMock.anyLong())).andReturn(
+                Boolean.FALSE);
+        EasyMock.expect(mMockMonitor.getDeviceState()).andReturn(TestDeviceState.NOT_AVAILABLE);
+        // expect reboot
+        EasyMock.expect(mMockRunUtil.runTimedCmd(EasyMock.anyLong(), EasyMock.eq("fastboot"),
+                EasyMock.eq("-s"), EasyMock.eq("serial"), EasyMock.eq("reboot-bootloader"))).
+                andReturn(new CommandResult(CommandStatus.SUCCESS));
+        EasyMock.expect(mMockMonitor.waitForDeviceNotAvailable(EasyMock.anyLong())).andReturn(
                 Boolean.TRUE);
+        EasyMock.expect(mMockMonitor.waitForDeviceBootloader(EasyMock.anyLong())).andReturn(
+                Boolean.TRUE).times(2);
         replayMocks();
         mRecovery.recoverDeviceBootloader(mMockMonitor);
         verifyMocks();
@@ -183,9 +185,30 @@
      */
     public void testRecoverDeviceBootloader_online() throws Exception {
         mMockRunUtil.sleep(EasyMock.anyLong());
-        mMockMonitor.waitForDeviceBootloaderStateUpdate();
-        EasyMock.expect(mMockMonitor.getDeviceState()).andReturn(TestDeviceState.ONLINE).
-                times(2);
+        EasyMock.expect(mMockMonitor.waitForDeviceBootloader(EasyMock.anyLong())).andReturn(
+                Boolean.FALSE);
+        EasyMock.expect(mMockMonitor.getDeviceState()).andReturn(TestDeviceState.ONLINE);
+        EasyMock.expect(mMockMonitor.waitForDeviceOnline()).andReturn(mMockDevice);
+        mMockDevice.reboot("bootloader");
+        EasyMock.expect(mMockMonitor.waitForDeviceBootloader(EasyMock.anyLong())).andReturn(
+                Boolean.TRUE);
+        replayMocks();
+        mRecovery.recoverDeviceBootloader(mMockMonitor);
+        verifyMocks();
+    }
+
+    /**
+     * Test {@link WaitDeviceRecovery#recoverDeviceBootloader(IDeviceStateMonitor)} when device is
+     * initially unavailable, then comes online when bootloader is expected
+     */
+    public void testRecoverDeviceBootloader_unavailable_online() throws Exception {
+        mMockRunUtil.sleep(EasyMock.anyLong());
+        EasyMock.expect(mMockMonitor.waitForDeviceBootloader(EasyMock.anyLong())).andReturn(
+                Boolean.FALSE);
+        EasyMock.expect(mMockMonitor.getDeviceState()).andReturn(TestDeviceState.NOT_AVAILABLE);
+        EasyMock.expect(mMockMonitor.waitForDeviceBootloader(EasyMock.anyLong())).andReturn(
+                Boolean.FALSE);
+        EasyMock.expect(mMockMonitor.getDeviceState()).andReturn(TestDeviceState.ONLINE);
         EasyMock.expect(mMockMonitor.waitForDeviceOnline()).andReturn(mMockDevice);
         mMockDevice.reboot("bootloader");
         EasyMock.expect(mMockMonitor.waitForDeviceBootloader(EasyMock.anyLong())).andReturn(
@@ -201,11 +224,9 @@
      */
     public void testRecoverDeviceBootloader_unavailable_failure() throws Exception {
         mMockRunUtil.sleep(EasyMock.anyLong());
-        mMockMonitor.waitForDeviceBootloaderStateUpdate();
-        EasyMock.expect(mMockMonitor.getDeviceState()).andReturn(TestDeviceState.NOT_AVAILABLE).
-                times(2);
-        EasyMock.expect(mMockMonitor.waitForDeviceBootloader(EasyMock.anyLong())).andReturn(
+        EasyMock.expect(mMockMonitor.waitForDeviceBootloader(EasyMock.anyLong())).andStubReturn(
                 Boolean.FALSE);
+        EasyMock.expect(mMockMonitor.getDeviceState()).andStubReturn(TestDeviceState.NOT_AVAILABLE);
         replayMocks();
         try {
             mRecovery.recoverDeviceBootloader(mMockMonitor);