Pass correct install user when creating sessions

This CL passes the correct install user when creating sessions in
StagingManager and effectively reverts ag/6901350.

Add a test to ensure b/129397974 is not regressed as described in
b/158222747#comment11.

(Cherry-picked from 838dea54caba0c6a17d90757211bc897ef77c963)

Fix: 129744602
Fix: 158222747
Test: atest MultiUserRollbackTest#testStagedRollback
Merged-In: I718992240aa76898ff9e4220ea6a769ee8cd61f5
Change-Id: I718992240aa76898ff9e4220ea6a769ee8cd61f5

4 files changed, 125 insertions, 0 deletions

diff --git a/tests/RollbackTest/Android.bp b/tests/RollbackTest/Android.bp
index 2be4ae6bb214..3ccbad84f2e0 100644
--- a/tests/RollbackTest/Android.bp
+++ b/tests/RollbackTest/Android.bp
@@ -48,6 +48,9 @@ java_test_host {
     name: "MultiUserRollbackTest",
     srcs: ["MultiUserRollbackTest/src/**/*.java"],
     libs: ["tradefed"],
+    static_libs: [
+        "frameworks-base-hostutils",
+    ],
     test_suites: ["general-tests"],
     test_config: "MultiUserRollbackTest.xml",
 }
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 42b886f0774f..f16084744853 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
@@ -24,6 +24,7 @@ import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -40,6 +41,9 @@ 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;
 
+    @Rule
+    public AbandonSessionsRule mHostTestRule = new AbandonSessionsRule(this);
+
     @After
     public void tearDown() throws Exception {
         removeSecondaryUserIfNecessary();
@@ -59,6 +63,30 @@ public class MultiUserRollbackTest extends BaseHostJUnit4Test {
         runPhaseForUsers("testBasic", mSecondaryUserId);
     }
 
+    /**
+     * Tests staged install/rollback works correctly on the 2nd user.
+     */
+    @Test
+    public void testStagedRollback() throws Exception {
+        runPhaseForUsers("testStagedRollback_Phase1", mSecondaryUserId);
+        getDevice().reboot();
+
+        // Need to unlock the user for device tests to run successfully
+        getDevice().startUser(mSecondaryUserId);
+        awaitUserUnlocked(mSecondaryUserId);
+        runPhaseForUsers("testStagedRollback_Phase2", mSecondaryUserId);
+        getDevice().reboot();
+
+        getDevice().startUser(mSecondaryUserId);
+        awaitUserUnlocked(mSecondaryUserId);
+        runPhaseForUsers("testStagedRollback_Phase3", mSecondaryUserId);
+        getDevice().reboot();
+
+        getDevice().startUser(mSecondaryUserId);
+        awaitUserUnlocked(mSecondaryUserId);
+        runPhaseForUsers("testStagedRollback_Phase4", mSecondaryUserId);
+    }
+
     @Test
     public void testMultipleUsers() throws Exception {
         runPhaseForUsers("testMultipleUsersInstallV1", mOriginalUserId, mSecondaryUserId);
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 8641f4d4013a..5d133a4de13d 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/MultiUserRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/MultiUserRollbackTest.java
@@ -115,4 +115,32 @@ public class MultiUserRollbackTest {
         assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
         InstallUtils.processUserData(TestApp.A);
     }
+
+    @Test
+    public void testStagedRollback_Phase1() throws Exception {
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(-1);
+        Install.single(TestApp.A1).setStaged().commit();
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(-1);
+    }
+
+    @Test
+    public void testStagedRollback_Phase2() throws Exception {
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
+        Install.single(TestApp.A2).setStaged().setEnableRollback().commit();
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
+    }
+
+    @Test
+    public void testStagedRollback_Phase3() throws Exception {
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(2);
+        RollbackInfo rollback = RollbackUtils.waitForAvailableRollback(TestApp.A);
+        assertThat(rollback).packagesContainsExactly(Rollback.from(TestApp.A2).to(TestApp.A1));
+        RollbackUtils.rollback(rollback.getRollbackId());
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(2);
+    }
+
+    @Test
+    public void testStagedRollback_Phase4() {
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
+    }
 }
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..5bfc75203ae3
--- /dev/null
+++ b/tests/utils/hostutils/src/com/android/tests/rollback/host/AbandonSessionsRule.java
@@ -0,0 +1,66 @@
+/*
+ * 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.ddmlib.Log;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
+
+import org.junit.rules.ExternalResource;
+
+public final class AbandonSessionsRule extends ExternalResource {
+    private static final String TAG = "AbandonSessionsRule";
+    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 e) {
+            mHost.getDevice().logOnDevice(TAG, Log.LogLevel.ERROR,
+                    "%s", "Failed to abandon sessions");
+        }
+    }
+
+    /**
+     * 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();
+    }
+}