diff --git a/tests/RollbackTest/Android.bp b/tests/RollbackTest/Android.bp
index 24623fba5bffd..2be4ae6bb2145 100644
--- a/tests/RollbackTest/Android.bp
+++ b/tests/RollbackTest/Android.bp
@@ -39,7 +39,7 @@ java_test_host {
     name: "NetworkStagedRollbackTest",
     srcs: ["NetworkStagedRollbackTest/src/**/*.java"],
     libs: ["tradefed"],
-    static_libs: ["testng", "RollbackTestLib"],
+    static_libs: ["RollbackTestLib"],
     test_suites: ["general-tests"],
     test_config: "NetworkStagedRollbackTest.xml",
 }
diff --git a/tests/RollbackTest/NetworkStagedRollbackTest/src/com/android/tests/rollback/host/NetworkStagedRollbackTest.java b/tests/RollbackTest/NetworkStagedRollbackTest/src/com/android/tests/rollback/host/NetworkStagedRollbackTest.java
index f6dcff49c80ce..57c52f9c30219 100644
--- a/tests/RollbackTest/NetworkStagedRollbackTest/src/com/android/tests/rollback/host/NetworkStagedRollbackTest.java
+++ b/tests/RollbackTest/NetworkStagedRollbackTest/src/com/android/tests/rollback/host/NetworkStagedRollbackTest.java
@@ -20,7 +20,6 @@ import static com.android.tests.rollback.host.WatchdogEventLogger.watchdogEventO
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
-import static org.testng.Assert.assertThrows;
 
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
@@ -31,6 +30,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Runs the network rollback tests.
@@ -83,11 +83,12 @@ public class NetworkStagedRollbackTest extends BaseHostJUnit4Test {
 
             // Verify rollback was enabled
             runPhase("testNetworkFailedRollback_Phase2");
-            assertThrows(AssertionError.class, () -> runPhase("testNetworkFailedRollback_Phase3"));
-
+            // 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("testNetworkFailedRollback_Phase4");
+            runPhase("testNetworkFailedRollback_Phase3");
 
             List<String> watchdogEvents = mLogger.getWatchdogLoggingEvents();
             assertTrue(watchdogEventOccurred(watchdogEvents, ROLLBACK_INITIATE, null,
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/NetworkStagedRollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/NetworkStagedRollbackTest.java
index 35bc65a24d630..7977b9a27f8fa 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/NetworkStagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/NetworkStagedRollbackTest.java
@@ -119,15 +119,6 @@ public class NetworkStagedRollbackTest {
 
     @Test
     public void testNetworkFailedRollback_Phase3() throws Exception {
-        // Sleep for > health check deadline (120s to trigger rollback + 120s to reboot)
-        // 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(240));
-    }
-
-    @Test
-    public void testNetworkFailedRollback_Phase4() throws Exception {
         RollbackManager rm = RollbackUtils.getRollbackManager();
         assertThat(getUniqueRollbackInfoForPackage(rm.getRecentlyCommittedRollbacks(),
                 getNetworkStackPackageName())).isNotNull();