From b13131f10c55a7b20fe3a558437445bb135ac94f Mon Sep 17 00:00:00 2001
From: JW Wang <wangchun@google.com>
Date: Mon, 30 Dec 2019 16:13:12 +0800
Subject: [PATCH] Clean up the rollback when test is done

When multiple staged sessions are installed on a device which doesn't
support checkpoint, only the 1st will prevail. The test failed when
TestApp.A is not the 1st rollback to commit.

We need to clean up the rollback left by testRollbackWhitelistedApp so
TestApp.A is always the only and first one to commit.

Bug: 146991129
Test: atest StagedRollbackTest
Change-Id: I4afdace934b7b7e3b198169dc2eb0da66c65eb94
---
 .../android/tests/rollback/StagedRollbackTest.java  | 11 +++++++++++
 .../tests/rollback/host/StagedRollbackTest.java     | 13 ++++++++++---
 2 files changed, 21 insertions(+), 3 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 879ac64c6a550..bb3906d20df69 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -185,6 +185,12 @@ public class StagedRollbackTest {
      */
     @Test
     public void testNativeWatchdogTriggersRollback_Phase1() throws Exception {
+        // When multiple staged sessions are installed on a device which doesn't support checkpoint,
+        // only the 1st one will prevail. We have to check no other rollbacks available to ensure
+        // TestApp.A is always the 1st and the only one to commit so rollback can work as intended.
+        // If there are leftover rollbacks from previous tests, this assertion will fail.
+        assertThat(RollbackUtils.getRollbackManager().getAvailableRollbacks()).isEmpty();
+
         Uninstall.packages(TestApp.A);
         Install.single(TestApp.A1).commit();
         assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
@@ -373,6 +379,11 @@ public class StagedRollbackTest {
         assertThat(RollbackUtils.getAvailableRollback(getModuleMetadataPackageName())).isNotNull();
     }
 
+    @Test
+    public void testRollbackWhitelistedApp_cleanUp() throws Exception {
+        RollbackUtils.getRollbackManager().expireRollbackForPackage(getModuleMetadataPackageName());
+    }
+
     private static void runShellCommand(String cmd) {
         ParcelFileDescriptor pfd = InstrumentationRegistry.getInstrumentation().getUiAutomation()
                 .executeShellCommand(cmd);
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 07d829d2d0bbe..8ab4c4ceda3a6 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
@@ -180,9 +180,16 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
      */
     @Test
     public void testRollbackWhitelistedApp() throws Exception {
-        runPhase("testRollbackWhitelistedApp_Phase1");
-        getDevice().reboot();
-        runPhase("testRollbackWhitelistedApp_Phase2");
+        try {
+            runPhase("testRollbackWhitelistedApp_Phase1");
+            getDevice().reboot();
+            runPhase("testRollbackWhitelistedApp_Phase2");
+        } finally {
+            // testNativeWatchdogTriggersRollback will fail if multiple staged sessions are
+            // committed on a device which doesn't support checkpoint. Let's clean up the rollback
+            // so there is only one rollback to commit when testing native crashes.
+            runPhase("testRollbackWhitelistedApp_cleanUp");
+        }
     }
 
     private void crashProcess(String processName, int numberOfCrashes) throws Exception {