1 files changed, 25 insertions, 34 deletions
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 db5e7f8d337f..cdc21ef8f971 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -158,18 +158,6 @@ public class StagedRollbackTest {
         assertThat(rollback.getCommittedSessionId()).isNotEqualTo(-1);
     }
 
-    @Test
-    public void resetNetworkStack() throws Exception {
-        RollbackManager rm = RollbackUtils.getRollbackManager();
-        String networkStack = getNetworkStackPackageName();
-
-        rm.expireRollbackForPackage(networkStack);
-        Uninstall.packages(networkStack);
-
-        assertThat(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(),
-                        networkStack)).isNull();
-    }
-
     /**
      * Stage install ModuleMetadata package to simulate a Mainline module update.
      */
@@ -210,43 +198,35 @@ public class StagedRollbackTest {
 
     @Test
     public void testNetworkFailedRollback_Phase1() throws Exception {
-        resetNetworkStack();
-    }
-
-    @Test
-    public void testNetworkFailedRollback_Phase2() throws Exception {
-        assertNetworkStackRollbackAvailable();
-    }
+        RollbackManager rm = RollbackUtils.getRollbackManager();
+        String networkStack = getNetworkStackPackageName();
 
-    @Test
-    public void testNetworkFailedRollback_Phase3() throws Exception {
-        assertNoNetworkStackRollbackCommitted();
-    }
+        rm.expireRollbackForPackage(networkStack);
+        Uninstall.packages(networkStack);
 
-    @Test
-    public void testNetworkFailedRollback_Phase4() throws Exception {
-        assertNetworkStackRollbackCommitted();
+        assertThat(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(),
+                        networkStack)).isNull();
     }
 
     @Test
-    public void assertNetworkStackRollbackAvailable() throws Exception {
+    public void testNetworkFailedRollback_Phase2() throws Exception {
         RollbackManager rm = RollbackUtils.getRollbackManager();
         assertThat(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(),
                         getNetworkStackPackageName())).isNotNull();
     }
 
     @Test
-    public void assertNetworkStackRollbackCommitted() throws Exception {
+    public void testNetworkFailedRollback_Phase3() throws Exception {
         RollbackManager rm = RollbackUtils.getRollbackManager();
         assertThat(getUniqueRollbackInfoForPackage(rm.getRecentlyCommittedRollbacks(),
-                        getNetworkStackPackageName())).isNotNull();
+                        getNetworkStackPackageName())).isNull();
     }
 
     @Test
-    public void assertNoNetworkStackRollbackCommitted() throws Exception {
+    public void testNetworkFailedRollback_Phase4() throws Exception {
         RollbackManager rm = RollbackUtils.getRollbackManager();
         assertThat(getUniqueRollbackInfoForPackage(rm.getRecentlyCommittedRollbacks(),
-                        getNetworkStackPackageName())).isNull();
+                        getNetworkStackPackageName())).isNotNull();
     }
 
     private String getNetworkStackPackageName() {
@@ -294,17 +274,28 @@ public class StagedRollbackTest {
 
     @Test
     public void testNetworkPassedDoesNotRollback_Phase1() throws Exception {
-        resetNetworkStack();
+        RollbackManager rm = RollbackUtils.getRollbackManager();
+        String networkStack = getNetworkStackPackageName();
+
+        rm.expireRollbackForPackage(networkStack);
+        Uninstall.packages(networkStack);
+
+        assertThat(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(),
+                        networkStack)).isNull();
     }
 
     @Test
     public void testNetworkPassedDoesNotRollback_Phase2() throws Exception {
-        assertNetworkStackRollbackAvailable();
+        RollbackManager rm = RollbackUtils.getRollbackManager();
+        assertThat(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(),
+                        getNetworkStackPackageName())).isNotNull();
     }
 
     @Test
     public void testNetworkPassedDoesNotRollback_Phase3() throws Exception {
-        assertNoNetworkStackRollbackCommitted();
+        RollbackManager rm = RollbackUtils.getRollbackManager();
+        assertThat(getUniqueRollbackInfoForPackage(rm.getRecentlyCommittedRollbacks(),
+                        getNetworkStackPackageName())).isNull();
     }
 
     @Nullable