Rewrite testNetworkFailedRollback_Phase3 (1/n)

Move test code to device side so we can assert reboot more reliably.

Bug: 143035398
Test: atest StagedRollbackTest
Change-Id: I84f65ce17ad38b524b26436933f9c090c1b581c5
diff --git a/tests/RollbackTest/Android.bp b/tests/RollbackTest/Android.bp
index 085c53c..2bc129a 100644
--- a/tests/RollbackTest/Android.bp
+++ b/tests/RollbackTest/Android.bp
@@ -25,6 +25,7 @@
     name: "StagedRollbackTest",
     srcs: ["StagedRollbackTest/src/**/*.java"],
     libs: ["tradefed"],
+    static_libs: ["testng"],
     test_suites: ["general-tests"],
     test_config: "StagedRollbackTest.xml",
 }
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
index b51aad1..aa09f05 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -229,13 +229,21 @@
         RollbackManager rm = RollbackUtils.getRollbackManager();
         assertThat(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(),
                         getNetworkStackPackageName())).isNotNull();
+
+        // Sleep for < health check deadline
+        Thread.sleep(TimeUnit.SECONDS.toMillis(5));
+        // Verify rollback was not executed before health check deadline
+        assertThat(getUniqueRollbackInfoForPackage(rm.getRecentlyCommittedRollbacks(),
+                        getNetworkStackPackageName())).isNull();
     }
 
     @Test
     public void testNetworkFailedRollback_Phase3() throws Exception {
-        RollbackManager rm = RollbackUtils.getRollbackManager();
-        assertThat(getUniqueRollbackInfoForPackage(rm.getRecentlyCommittedRollbacks(),
-                        getNetworkStackPackageName())).isNull();
+        // Sleep for > health check deadline
+        // The device is expected to reboot during sleeping. This device method will fail and
+        // the host will catch the assertion. If reboot doesn't happen, the host will fail the
+        // assertion.
+        Thread.sleep(TimeUnit.SECONDS.toMillis(260));
     }
 
     @Test
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 f7fe6c7..afbe9156 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
@@ -17,6 +17,7 @@
 package com.android.tests.rollback.host;
 
 import static org.junit.Assert.assertTrue;
+import static org.testng.Assert.assertThrows;
 
 import com.android.ddmlib.Log.LogLevel;
 import com.android.tradefed.log.LogUtil.CLog;
@@ -136,22 +137,7 @@
 
         // Verify rollback was enabled
         runPhase("testNetworkFailedRollback_Phase2");
-
-        // Sleep for < health check deadline
-        Thread.sleep(5000);
-        // Verify rollback was not executed before health check deadline
-        runPhase("testNetworkFailedRollback_Phase3");
-        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 testNetworkFailedRollback_Phase4 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.
-        }
+        assertThrows(AssertionError.class, () -> runPhase("testNetworkFailedRollback_Phase3"));
 
         getDevice().waitForDeviceAvailable();
         // Verify rollback was executed after health check deadline