Merge "Fix a flaky test"

tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java

@@ -128,7 +128,7 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
         runPhase("resetNetworkStack");
         // Reduce health check deadline
         getDevice().executeShellCommand("device_config put rollback "
-                + "watchdog_request_timeout_millis 300000");
+                + "watchdog_request_timeout_millis 120000");
         // Simulate re-installation of new NetworkStack with rollbacks enabled
         getDevice().executeShellCommand("pm install -r --staged --enable-rollback "
                 + getNetworkStackPath());
@@ -146,18 +146,10 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
         Thread.sleep(5000);
         // Verify rollback was not executed before health check deadline
         runPhase("assertNoNetworkStackRollbackCommitted");
-        try {
-            // This is expected to fail due to the device being rebooted out
-            // from underneath the test. If this fails for reasons other than
-            // the device reboot, those failures should result in failure of
-            // the assertNetworkStackExecutedRollback phase.
-            CLog.logAndDisplay(LogLevel.INFO, "Sleep and expect to fail while sleeping");
-            // Sleep for > health check deadline
-            Thread.sleep(260000);
-        } catch (AssertionError e) {
-            // AssertionError is expected.
-        }
 
+        // Wait for reboot to happen
+        assertTrue(getDevice().waitForDeviceNotAvailable(TimeUnit.MINUTES.toMillis(5)));
+        // Wait for reboot to complete and device to become available
         getDevice().waitForDeviceAvailable();
         // Verify rollback was executed after health check deadline
         runPhase("assertNetworkStackRollbackCommitted");