Rewrite how we abandon sessions

The original code is flawed in that `pm install-abandon` only abandons
the 1st session returned by `pm get-stagedsessions ...`.

1. move AbandonSessionsRule to be shared by multiple host tests
2. use AbandonSessionsRule to do the job

Bug: 160754072
Bug: 185973415
Test: StagedInstallInternalTest
Change-Id: Ib7b32fbd7b1133ac6a8e6782234a4fe2c5a782bd
Merged-In: Ib7b32fbd7b1133ac6a8e6782234a4fe2c5a782bd
(cherry picked from commit 3ac333f2d20e240093235d957ba67ce174fe2db8)

4 files changed, 78 insertions, 14 deletions

diff --git a/tests/RollbackTest/Android.bp b/tests/RollbackTest/Android.bp
index 2be4ae6bb214..a23df920b396 100644
--- a/tests/RollbackTest/Android.bp
+++ b/tests/RollbackTest/Android.bp
@@ -29,7 +29,12 @@ java_test_host {
     name: "StagedRollbackTest",
     srcs: ["StagedRollbackTest/src/**/*.java"],
     libs: ["tradefed"],
-    static_libs: ["testng", "compatibility-tradefed", "RollbackTestLib"],
+    static_libs: [
+        "compatibility-tradefed",
+        "frameworks-base-hostutils",
+        "RollbackTestLib",
+        "testng",
+    ],
     test_suites: ["general-tests"],
     test_config: "StagedRollbackTest.xml",
     data: [":com.android.apex.apkrollback.test_v1"],
@@ -39,7 +44,7 @@ java_test_host {
     name: "NetworkStagedRollbackTest",
     srcs: ["NetworkStagedRollbackTest/src/**/*.java"],
     libs: ["tradefed"],
-    static_libs: ["RollbackTestLib"],
+    static_libs: ["RollbackTestLib", "frameworks-base-hostutils"],
     test_suites: ["general-tests"],
     test_config: "NetworkStagedRollbackTest.xml",
 }
diff --git a/tests/StagedInstallTest/Android.bp b/tests/StagedInstallTest/Android.bp
index c3fdd695c2b7..65d15a7dccaa 100644
--- a/tests/StagedInstallTest/Android.bp
+++ b/tests/StagedInstallTest/Android.bp
@@ -24,7 +24,7 @@ java_test_host {
     name: "StagedInstallInternalTest",
     srcs: ["src/**/*.java"],
     libs: ["tradefed"],
-    static_libs: ["testng", "compatibility-tradefed"],
+    static_libs: ["testng", "compatibility-tradefed", "frameworks-base-hostutils"],
     test_suites: ["general-tests"],
     test_config: "StagedInstallInternalTest.xml",
 }
diff --git a/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java b/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
index 9b432f7d0ca5..530f885e62bf 100644
--- a/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
+++ b/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
@@ -21,12 +21,14 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertTrue;
 
 import com.android.ddmlib.Log;
+import com.android.tests.rollback.host.AbandonSessionsRule;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 import com.android.tradefed.util.ProcessInfo;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -35,7 +37,9 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
 
     private static final String TAG = StagedInstallInternalTest.class.getSimpleName();
     private static final long SYSTEM_SERVER_TIMEOUT_MS = 60 * 1000;
-    private boolean mWasRoot = false;
+
+    @Rule
+    public AbandonSessionsRule mHostTestRule = new AbandonSessionsRule(this);
 
     /**
      * Runs the given phase of a test by calling into the device.
@@ -62,21 +66,11 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
 
     @Before
     public void setUp() throws Exception {
-        mWasRoot = getDevice().isAdbRoot();
-        if (!mWasRoot) {
-            getDevice().enableAdbRoot();
-        }
         cleanUp();
-        // Abandon all staged sessions
-        getDevice().executeShellCommand("pm install-abandon $(pm get-stagedsessions --only-ready "
-                + "--only-parent --only-sessionid)");
     }
 
     @After
     public void tearDown() throws Exception {
-        if (!mWasRoot) {
-            getDevice().disableAdbRoot();
-        }
         cleanUp();
     }
 
@@ -90,7 +84,10 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
     private void restartSystemServer() throws Exception {
         // Restart the system server
         long oldStartTime = getDevice().getProcessByName("system_server").getStartTime();
+
+        getDevice().enableAdbRoot(); // Need root to restart system server
         assertThat(getDevice().executeShellCommand("am restart")).contains("Restart the system");
+        getDevice().disableAdbRoot();
 
         // Wait for new system server process to start
         long start = System.currentTimeMillis();
diff --git a/tests/utils/hostutils/src/com/android/tests/rollback/host/AbandonSessionsRule.java b/tests/utils/hostutils/src/com/android/tests/rollback/host/AbandonSessionsRule.java
new file mode 100644
index 000000000000..b08621314ee0
--- /dev/null
+++ b/tests/utils/hostutils/src/com/android/tests/rollback/host/AbandonSessionsRule.java
@@ -0,0 +1,62 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.tests.rollback.host;
+
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
+
+import org.junit.rules.ExternalResource;
+
+public class AbandonSessionsRule extends ExternalResource {
+    private final BaseHostJUnit4Test mHost;
+
+    public AbandonSessionsRule(BaseHostJUnit4Test host) {
+        mHost = host;
+    }
+
+    @Override
+    protected void before() throws Throwable {
+        abandonSessions(mHost.getDevice());
+    }
+
+    @Override
+    protected void after() {
+        try {
+            abandonSessions(mHost.getDevice());
+        } catch (Exception ignore) {
+        }
+    }
+
+    /**
+     * Abandons all sessions to prevent interference in our tests.
+     */
+    private static void abandonSessions(ITestDevice device) throws Exception {
+        // No point in abandoning applied or failed sessions. We care about ready sessions only.
+        String cmdListReadySessions =
+                "pm list staged-sessions --only-sessionid --only-parent --only-ready";
+        String output = device.executeShellCommand(cmdListReadySessions);
+        if (output.trim().isEmpty()) {
+            // No sessions to abandon
+            return;
+        }
+        // Ensure we have sufficient privilege to abandon sessions from other apps
+        device.enableAdbRoot();
+        device.executeShellCommand("for i in $(" + cmdListReadySessions
+                + "); do pm install-abandon $i; done");
+        device.disableAdbRoot();
+    }
+}