Fix a flaky test

See b/152840581#comment2 for the analysis.

1. call #waitForDeviceNotAvailable to wait for reboot reliably.
2. reduce watchdog timeout to 120s which works well for me and speeds up
   the test.

Bug: 152840581
Test: atest StagedRollbackTest
Merged-In: Ibbc9998c74681b1b8fc29b90ea66406f39ae7443
Change-Id: Ic791fbc00c231a5f24751b36e99515cfad13b7e7

1 files changed, 4 insertions, 12 deletions

diff --git a/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java b/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
index a590fc4e28d5..47d01734be40 100644
--- a/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
+++ b/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");