Make rollback-app support --staged-ready-timeout flag

Staged rollback should be waiting for pre-reboot verificatio to complete
by default.

Also added a test to ensure it's working properly.

Bug: 162958790
Test: atest StagedInstallInternalTest#testAdbRollbackAppWaitsForStagedReady
Change-Id: I5dcaabaa2ecf9a5e137773821dd5920a09629aeb

services/core/java/com/android/server/pm/PackageManagerShellCommand.java

@@ -464,9 +464,20 @@ class PackageManagerShellCommand extends ShellCommand {
         return 1;
     }
 
-    private int runRollbackApp() {
+    private int runRollbackApp() throws RemoteException {
         final PrintWriter pw = getOutPrintWriter();
 
+        String opt;
+        long stagedReadyTimeoutMs = DEFAULT_STAGED_READY_TIMEOUT_MS;
+        while ((opt = getNextOption()) != null) {
+            switch (opt) {
+                case "--staged-ready-timeout":
+                    stagedReadyTimeoutMs = Long.parseLong(getNextArgRequired());
+                    break;
+                default:
+                    throw new IllegalArgumentException("Unknown option: " + opt);
+            }
+        }
         final String packageName = getNextArgRequired();
         if (packageName == null) {
             pw.println("Error: package name not specified");
@@ -496,14 +507,21 @@ class PackageManagerShellCommand extends ShellCommand {
         final Intent result = receiver.getResult();
         final int status = result.getIntExtra(RollbackManager.EXTRA_STATUS,
                 RollbackManager.STATUS_FAILURE);
-        if (status == RollbackManager.STATUS_SUCCESS) {
-            pw.println("Success");
-            return 0;
-        } else {
+
+        if (status != RollbackManager.STATUS_SUCCESS) {
             pw.println("Failure ["
                     + result.getStringExtra(RollbackManager.EXTRA_STATUS_MESSAGE) + "]");
             return 1;
         }
+
+        if (rollback.isStaged() && stagedReadyTimeoutMs > 0) {
+            final int committedSessionId = rollback.getCommittedSessionId();
+            return doWaitForStagedSessionReady(committedSessionId, stagedReadyTimeoutMs, pw);
+        }
+
+        pw.println("Success");
+        return 0;
+
     }
 
     private void setParamsSize(InstallParams params, List<String> inPaths) {
@@ -1309,7 +1327,7 @@ class PackageManagerShellCommand extends ShellCommand {
             abandonSession = false;
 
             if (params.sessionParams.isStaged && params.stagedReadyTimeoutMs > 0) {
-                return doWaitForStagedSessionRead(sessionId, params.stagedReadyTimeoutMs, pw);
+                return doWaitForStagedSessionReady(sessionId, params.stagedReadyTimeoutMs, pw);
             }
 
             pw.println("Success");
@@ -1324,7 +1342,7 @@ class PackageManagerShellCommand extends ShellCommand {
         }
     }
 
-    private int doWaitForStagedSessionRead(int sessionId, long timeoutMs, PrintWriter pw)
+    private int doWaitForStagedSessionReady(int sessionId, long timeoutMs, PrintWriter pw)
               throws RemoteException {
         Preconditions.checkArgument(timeoutMs > 0);
         PackageInstaller.SessionInfo si = mInterface.getPackageInstaller()
@@ -1393,7 +1411,7 @@ class PackageManagerShellCommand extends ShellCommand {
         final PackageInstaller.SessionInfo si = mInterface.getPackageInstaller()
                 .getSessionInfo(sessionId);
         if (si != null && si.isStaged() && stagedReadyTimeoutMs > 0) {
-            return doWaitForStagedSessionRead(sessionId, stagedReadyTimeoutMs, pw);
+            return doWaitForStagedSessionReady(sessionId, stagedReadyTimeoutMs, pw);
         }
         pw.println("Success");
         return 0;

tests/StagedInstallTest/Android.bp

@@ -26,7 +26,7 @@ android_test_helper_app {
 java_test_host {
     name: "StagedInstallInternalTest",
     srcs: ["src/**/*.java"],
-    libs: ["tradefed"],
+    libs: ["tradefed", "cts-shim-host-lib"],
     static_libs: [
         "testng",
         "compatibility-tradefed",

tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java

@@ -16,6 +16,8 @@
 
 package com.android.tests.stagedinstallinternal.host;
 
+import static com.android.cts.shim.lib.ShimPackage.SHIM_APEX_PACKAGE_NAME;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertTrue;
@@ -197,6 +199,24 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
         assertThat(sessionId).isEmpty();
     }
 
+    // Test rollback-app command waits for staged sessions to be ready
+    @Test
+    public void testAdbRollbackAppWaitsForStagedReady() throws Exception {
+        assumeTrue("Device does not support updating APEX",
+                mHostUtils.isApexUpdateSupported());
+
+        final File apexFile = mTestUtils.getTestFile(SHIM_V2);
+        String output = getDevice().executeAdbCommand("install", "--staged",
+                "--enable-rollback", apexFile.getAbsolutePath());
+        assertThat(output).contains("Reboot device to apply staged session");
+        getDevice().reboot();
+        output = getDevice().executeShellCommand("pm rollback-app " + SHIM_APEX_PACKAGE_NAME);
+        assertThat(output).contains("Reboot device to apply staged session");
+        final String sessionId = getDevice().executeShellCommand(
+                "pm get-stagedsessions --only-ready --only-parent --only-sessionid").trim();
+        assertThat(sessionId).isNotEmpty();
+    }
+
     @Test
     public void testAdbInstallMultiPackageCommandWorks() throws Exception {
         assumeTrue("Device does not support updating APEX",