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 8f36b6e95c327..06be5e2a0b9d8 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -251,11 +251,11 @@ public class StagedRollbackTest {
 
     @Test
     public void testNetworkFailedRollback_Phase3() throws Exception {
-        // Sleep for > health check deadline
+        // 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(120));
+        Thread.sleep(TimeUnit.SECONDS.toMillis(240));
     }
 
     @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 bfd5dd5a33def..557d498ce8eec 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
@@ -22,6 +22,7 @@ import static org.testng.Assert.assertThrows;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 
+import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -47,6 +48,11 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
                     phase));
     }
 
+    @Before
+    public void setUp() throws Exception {
+        getDevice().reboot();
+    }
+
     /**
      * Tests watchdog triggered staged rollbacks involving only apks.
      */