From 864d59ed70478638cb54d92de5c8dc45c60a488b Mon Sep 17 00:00:00 2001
From: JW Wang <wangchun@google.com>
Date: Wed, 3 Jun 2020 17:33:45 +0800
Subject: [PATCH] Clean up rollbacks and test apps

See comment#6.

Clean up rollbacks and test apps before and after running tests.
This ensure we are in a good and expected condition when running tests.

Bug: 154080003
Test: atest RollbackTest
Test: atest MultiUserRollbackTest
Test: atest NetworkStagedRollbackTest
Test: atest StagedRollbackTest

Change-Id: If10fc8865e644611f1cf6a9db9d6dca6c7f11aee
---
 tests/RollbackTest/MultiUserRollbackTest.xml  |  6 ++++++
 .../rollback/host/MultiUserRollbackTest.java  |  9 ++-------
 .../host/NetworkStagedRollbackTest.java       |  2 ++
 .../tests/rollback/MultiUserRollbackTest.java | 12 +++++++++++
 .../rollback/NetworkStagedRollbackTest.java   | 20 +++++++++++++------
 .../android/tests/rollback/RollbackTest.java  | 17 ++++++++++++++++
 .../tests/rollback/StagedRollbackTest.java    |  5 ++++-
 tests/RollbackTest/StagedRollbackTest.xml     |  6 ++++++
 8 files changed, 63 insertions(+), 14 deletions(-)

diff --git a/tests/RollbackTest/MultiUserRollbackTest.xml b/tests/RollbackTest/MultiUserRollbackTest.xml
index ba86c3ff6777d..2f62af1856da4 100644
--- a/tests/RollbackTest/MultiUserRollbackTest.xml
+++ b/tests/RollbackTest/MultiUserRollbackTest.xml
@@ -15,6 +15,12 @@
 -->
 <configuration description="Runs rollback tests for multiple users">
     <option name="test-suite-tag" value="MultiUserRollbackTest" />
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="run-command" value="pm uninstall com.android.cts.install.lib.testapp.A" />
+        <option name="run-command" value="pm uninstall com.android.cts.install.lib.testapp.B" />
+        <option name="teardown-command" value="pm uninstall com.android.cts.install.lib.testapp.A" />
+        <option name="teardown-command" value="pm uninstall com.android.cts.install.lib.testapp.B" />
+    </target_preparer>
     <test class="com.android.tradefed.testtype.HostTest" >
         <option name="class" value="com.android.tests.rollback.host.MultiUserRollbackTest" />
     </test>
diff --git a/tests/RollbackTest/MultiUserRollbackTest/src/com/android/tests/rollback/host/MultiUserRollbackTest.java b/tests/RollbackTest/MultiUserRollbackTest/src/com/android/tests/rollback/host/MultiUserRollbackTest.java
index a4c81d5775223..42b886f0774f7 100644
--- a/tests/RollbackTest/MultiUserRollbackTest/src/com/android/tests/rollback/host/MultiUserRollbackTest.java
+++ b/tests/RollbackTest/MultiUserRollbackTest/src/com/android/tests/rollback/host/MultiUserRollbackTest.java
@@ -40,23 +40,18 @@ public class MultiUserRollbackTest extends BaseHostJUnit4Test {
     private static final long SWITCH_USER_COMPLETED_NUMBER_OF_POLLS = 60;
     private static final long SWITCH_USER_COMPLETED_POLL_INTERVAL_IN_MILLIS = 1000;
 
-    private void cleanUp() throws Exception {
-        getDevice().executeShellCommand("pm rollback-app com.android.cts.install.lib.testapp.A");
-        getDevice().executeShellCommand("pm uninstall com.android.cts.install.lib.testapp.A");
-    }
-
     @After
     public void tearDown() throws Exception {
-        cleanUp();
         removeSecondaryUserIfNecessary();
+        runPhaseForUsers("cleanUp", mOriginalUserId);
     }
 
     @Before
     public void setup() throws Exception {
-        cleanUp();
         mOriginalUserId = getDevice().getCurrentUser();
         createAndStartSecondaryUser();
         installPackage("RollbackTest.apk", "--user all");
+        runPhaseForUsers("cleanUp", mOriginalUserId);
     }
 
     @Test
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 57c52f9c30219..61d7c763e8d7c 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
@@ -59,12 +59,14 @@ public class NetworkStagedRollbackTest extends BaseHostJUnit4Test {
 
     @Before
     public void setUp() throws Exception {
+        runPhase("cleanUp");
         mLogger.start(getDevice());
     }
 
     @After
     public void tearDown() throws Exception {
         mLogger.stop();
+        runPhase("cleanUp");
     }
 
     /**
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/MultiUserRollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/MultiUserRollbackTest.java
index 400bb04f0fabf..8641f4d4013ab 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/MultiUserRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/MultiUserRollbackTest.java
@@ -22,6 +22,7 @@ import static com.google.common.truth.Truth.assertThat;
 
 import android.Manifest;
 import android.content.rollback.RollbackInfo;
+import android.content.rollback.RollbackManager;
 
 import com.android.cts.install.lib.Install;
 import com.android.cts.install.lib.InstallUtils;
@@ -53,6 +54,17 @@ public class MultiUserRollbackTest {
         InstallUtils.dropShellPermissionIdentity();
     }
 
+    @Test
+    public void cleanUp() {
+        RollbackManager rm = RollbackUtils.getRollbackManager();
+        rm.getAvailableRollbacks().stream().flatMap(info -> info.getPackages().stream())
+                .map(info -> info.getPackageName()).forEach(rm::expireRollbackForPackage);
+        rm.getRecentlyCommittedRollbacks().stream().flatMap(info -> info.getPackages().stream())
+                .map(info -> info.getPackageName()).forEach(rm::expireRollbackForPackage);
+        assertThat(rm.getAvailableRollbacks()).isEmpty();
+        assertThat(rm.getRecentlyCommittedRollbacks()).isEmpty();
+    }
+
     @Test
     public void testBasic() throws Exception {
         new RollbackTest().testBasic();
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 8fb59c7c3e2c9..42b0c608822ec 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/NetworkStagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/NetworkStagedRollbackTest.java
@@ -19,6 +19,8 @@ package com.android.tests.rollback;
 import static com.android.cts.rollback.lib.RollbackInfoSubject.assertThat;
 import static com.android.cts.rollback.lib.RollbackUtils.getUniqueRollbackInfoForPackage;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import android.Manifest;
 import android.content.ComponentName;
 import android.content.Intent;
@@ -90,15 +92,24 @@ public class NetworkStagedRollbackTest {
         InstallUtils.dropShellPermissionIdentity();
     }
 
+    @Test
+    public void cleanUp() {
+        RollbackManager rm = RollbackUtils.getRollbackManager();
+        rm.getAvailableRollbacks().stream().flatMap(info -> info.getPackages().stream())
+                .map(info -> info.getPackageName()).forEach(rm::expireRollbackForPackage);
+        rm.getRecentlyCommittedRollbacks().stream().flatMap(info -> info.getPackages().stream())
+                .map(info -> info.getPackageName()).forEach(rm::expireRollbackForPackage);
+        assertThat(rm.getAvailableRollbacks()).isEmpty();
+        assertThat(rm.getRecentlyCommittedRollbacks()).isEmpty();
+        uninstallNetworkStackPackage();
+    }
+
     @Test
     public void testNetworkFailedRollback_Phase1() throws Exception {
         // Remove available rollbacks and uninstall NetworkStack on /data/
         RollbackManager rm = RollbackUtils.getRollbackManager();
         String networkStack = getNetworkStackPackageName();
 
-        rm.expireRollbackForPackage(networkStack);
-        uninstallNetworkStackPackage();
-
         assertThat(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(),
                 networkStack)).isNull();
 
@@ -153,9 +164,6 @@ public class NetworkStagedRollbackTest {
         RollbackManager rm = RollbackUtils.getRollbackManager();
         String networkStack = getNetworkStackPackageName();
 
-        rm.expireRollbackForPackage(networkStack);
-        uninstallNetworkStackPackage();
-
         assertThat(getUniqueRollbackInfoForPackage(rm.getAvailableRollbacks(),
                 networkStack)).isNull();
 
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
index 48b5bed609d1d..dd08771b220db 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
@@ -48,6 +48,8 @@ import com.android.cts.rollback.lib.Rollback;
 import com.android.cts.rollback.lib.RollbackBroadcastReceiver;
 import com.android.cts.rollback.lib.RollbackUtils;
 
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -81,6 +83,21 @@ public class RollbackTest {
                         pri -> packageName.equals(pri.getPackageName())));
     }
 
+    @Before
+    @After
+    public void cleanUp() {
+        try {
+            InstallUtils.adoptShellPermissionIdentity(Manifest.permission.TEST_MANAGE_ROLLBACKS);
+            RollbackManager rm = RollbackUtils.getRollbackManager();
+            rm.getAvailableRollbacks().stream().flatMap(info -> info.getPackages().stream())
+                    .map(info -> info.getPackageName()).forEach(rm::expireRollbackForPackage);
+            rm.getRecentlyCommittedRollbacks().stream().flatMap(info -> info.getPackages().stream())
+                    .map(info -> info.getPackageName()).forEach(rm::expireRollbackForPackage);
+        } finally {
+            InstallUtils.dropShellPermissionIdentity();
+        }
+    }
+
     /**
      * Test basic rollbacks.
      */
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 6c9ffe2a7facb..00bd4cf388ced 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -360,7 +360,10 @@ public class StagedRollbackTest {
         RollbackManager rm = RollbackUtils.getRollbackManager();
         rm.getAvailableRollbacks().stream().flatMap(info -> info.getPackages().stream())
                 .map(info -> info.getPackageName()).forEach(rm::expireRollbackForPackage);
-        assertThat(RollbackUtils.getRollbackManager().getAvailableRollbacks()).isEmpty();
+        rm.getRecentlyCommittedRollbacks().stream().flatMap(info -> info.getPackages().stream())
+                .map(info -> info.getPackageName()).forEach(rm::expireRollbackForPackage);
+        assertThat(rm.getAvailableRollbacks()).isEmpty();
+        assertThat(rm.getRecentlyCommittedRollbacks()).isEmpty();
     }
 
     private static final String APK_IN_APEX_TESTAPEX_NAME = "com.android.apex.apkrollback.test";
diff --git a/tests/RollbackTest/StagedRollbackTest.xml b/tests/RollbackTest/StagedRollbackTest.xml
index 2750d3765c20b..83fef8e0a04ba 100644
--- a/tests/RollbackTest/StagedRollbackTest.xml
+++ b/tests/RollbackTest/StagedRollbackTest.xml
@@ -19,6 +19,12 @@
         <option name="cleanup-apks" value="true" />
         <option name="test-file-name" value="RollbackTest.apk" />
     </target_preparer>
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="run-command" value="pm uninstall com.android.cts.install.lib.testapp.A" />
+        <option name="run-command" value="pm uninstall com.android.cts.install.lib.testapp.B" />
+        <option name="teardown-command" value="pm uninstall com.android.cts.install.lib.testapp.A" />
+        <option name="teardown-command" value="pm uninstall com.android.cts.install.lib.testapp.B" />
+    </target_preparer>
     <test class="com.android.tradefed.testtype.HostTest" >
         <option name="class" value="com.android.tests.rollback.host.StagedRollbackTest" />
     </test>