4 files changed, 17 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 1e286bb15c49..30f3f185203c 100644
--- a/tests/StagedInstallTest/Android.bp
+++ b/tests/StagedInstallTest/Android.bp
@@ -28,6 +28,7 @@ java_test_host {
         "testng",
         "compatibility-tradefed",
         "module_test_util",
+        "frameworks-base-hostutils",
     ],
     data: [
         ":com.android.apex.cts.shim.v2_prebuilt",
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 86d5fd80c108..e259c9efbfcd 100644
--- a/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
+++ b/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
@@ -22,6 +22,7 @@ import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeTrue;
 
 import com.android.ddmlib.Log;
+import com.android.tests.rollback.host.AbandonSessionsRule;
 import com.android.tests.util.ModuleTestUtils;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
@@ -29,6 +30,7 @@ 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;
 
@@ -39,7 +41,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);
 
     private static final String SHIM_V2 = "com.android.apex.cts.shim.v2.apex";
     private static final String APK_A = "TestAppAv1.apk";
@@ -71,21 +75,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();
     }
 
@@ -150,8 +144,11 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
 
     private void restartSystemServer() throws Exception {
         // Restart the system server
-        long oldStartTime = getDevice().getProcessByName("system_server").getStartTime();
+        final 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/RollbackTest/lib/src/com/android/tests/rollback/host/AbandonSessionsRule.java b/tests/utils/hostutils/src/com/android/tests/rollback/host/AbandonSessionsRule.java
index b08621314ee0..b08621314ee0 100644
--- a/tests/RollbackTest/lib/src/com/android/tests/rollback/host/AbandonSessionsRule.java
+++ b/tests/utils/hostutils/src/com/android/tests/rollback/host/AbandonSessionsRule.java