diff --git a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
index 09682c6..39b3203 100644
--- a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
+++ b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
@@ -1377,7 +1377,7 @@
         long timeoutMs = -1;
         while ((opt = getNextOption()) != null) {
             switch (opt) {
-                case "--wait":
+                case "--wait-for-staged-ready":
                     waitForStagedSessionReady = true;
                     // If there is only one remaining argument, then it represents the sessionId, we
                     // shouldn't try to parse it as timeoutMs.
@@ -2865,7 +2865,7 @@
                     }
                     sessionParams.installFlags |= PackageManager.INSTALL_ENABLE_ROLLBACK;
                     break;
-                case "--wait":
+                case "--wait-for-staged-ready":
                     params.mWaitForStagedSessionReady = true;
                     try {
                         params.timeoutMs = Long.parseLong(peekNextArg());
@@ -3597,7 +3597,7 @@
         pw.println("       [--preload] [--instant] [--full] [--dont-kill]");
         pw.println("       [--enable-rollback]");
         pw.println("       [--force-uuid internal|UUID] [--pkg PACKAGE] [-S BYTES]");
-        pw.println("       [--apex] [--wait TIMEOUT]");
+        pw.println("       [--apex] [--wait-for-staged-ready TIMEOUT]");
         pw.println("       [PATH [SPLIT...]|-]");
         pw.println("    Install an application.  Must provide the apk data to install, either as");
         pw.println("    file path(s) or '-' to read from stdin.  Options are:");
@@ -3625,8 +3625,8 @@
         pw.println("          3=device setup, 4=user request");
         pw.println("      --force-uuid: force install on to disk volume with given UUID");
         pw.println("      --apex: install an .apex file, not an .apk");
-        pw.println("      --wait: when performing staged install, wait TIMEOUT milliseconds");
-        pw.println("          for pre-reboot verification to complete. If TIMEOUT is not");
+        pw.println("      --wait-for-staged-ready: when performing staged install, wait TIMEOUT");
+        pw.println("          ms for pre-reboot verification to complete. If TIMEOUT is not");
         pw.println("          specified it will wait for " + DEFAULT_WAIT_MS + " milliseconds.");
         pw.println("");
         pw.println("  install-existing [--user USER_ID|all|current]");
diff --git a/tests/StagedInstallTest/Android.bp b/tests/StagedInstallTest/Android.bp
index c3fdd69..1e286bb 100644
--- a/tests/StagedInstallTest/Android.bp
+++ b/tests/StagedInstallTest/Android.bp
@@ -24,7 +24,15 @@
     name: "StagedInstallInternalTest",
     srcs: ["src/**/*.java"],
     libs: ["tradefed"],
-    static_libs: ["testng", "compatibility-tradefed"],
+    static_libs: [
+        "testng",
+        "compatibility-tradefed",
+        "module_test_util",
+    ],
+    data: [
+        ":com.android.apex.cts.shim.v2_prebuilt",
+        ":TestAppAv1",
+    ],
     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 9b432f7..86d5fd8 100644
--- a/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
+++ b/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
@@ -19,8 +19,10 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeTrue;
 
 import com.android.ddmlib.Log;
+import com.android.tests.util.ModuleTestUtils;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 import com.android.tradefed.util.ProcessInfo;
@@ -30,6 +32,8 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.io.File;
+
 @RunWith(DeviceJUnit4ClassRunner.class)
 public class StagedInstallInternalTest extends BaseHostJUnit4Test {
 
@@ -37,6 +41,11 @@
     private static final long SYSTEM_SERVER_TIMEOUT_MS = 60 * 1000;
     private boolean mWasRoot = false;
 
+    private static final String SHIM_V2 = "com.android.apex.cts.shim.v2.apex";
+    private static final String APK_A = "TestAppAv1.apk";
+
+    private final ModuleTestUtils mTestUtils = new ModuleTestUtils(this);
+
     /**
      * Runs the given phase of a test by calling into the device.
      * Throws an exception if the test phase fails.
@@ -87,6 +96,58 @@
         runPhase("testSystemServerRestartDoesNotAffectStagedSessions_Verify");
     }
 
+    @Test
+    public void testAdbStagedInstallWaitForReadyFlagWorks() throws Exception {
+        assumeTrue("Device does not support updating APEX",
+                mTestUtils.isApexUpdateSupported());
+
+        File apexFile = mTestUtils.getTestFile(SHIM_V2);
+        String output = getDevice().executeAdbCommand("install", "--staged",
+                "--wait-for-staged-ready", "60000", apexFile.getAbsolutePath());
+        assertThat(output).contains("Reboot device to apply staged session");
+        String sessionId = getDevice().executeShellCommand(
+                "pm get-stagedsessions --only-ready --only-parent --only-sessionid").trim();
+        assertThat(sessionId).isNotEmpty();
+    }
+
+    @Test
+    public void testAdbStagedInstallNoWaitFlagWorks() throws Exception {
+        assumeTrue("Device does not support updating APEX",
+                mTestUtils.isApexUpdateSupported());
+
+        File apexFile = mTestUtils.getTestFile(SHIM_V2);
+        String output = getDevice().executeAdbCommand("install", "--staged",
+                "--no-wait", apexFile.getAbsolutePath());
+        assertThat(output).doesNotContain("Reboot device to apply staged session");
+        assertThat(output).contains("Success");
+        String sessionId = getDevice().executeShellCommand(
+                "pm get-stagedsessions --only-ready --only-parent --only-sessionid").trim();
+        assertThat(sessionId).isEmpty();
+    }
+
+    @Test
+    public void testAdbInstallMultiPackageCommandWorks() throws Exception {
+        assumeTrue("Device does not support updating APEX",
+                mTestUtils.isApexUpdateSupported());
+
+        File apexFile = mTestUtils.getTestFile(SHIM_V2);
+        File apkFile = mTestUtils.getTestFile(APK_A);
+        String output = getDevice().executeAdbCommand("install-multi-package",
+                apexFile.getAbsolutePath(), apkFile.getAbsolutePath());
+        assertThat(output).contains("Created parent session");
+        assertThat(output).contains("Created child session");
+        assertThat(output).contains("Success. Reboot device to apply staged session");
+
+        // Ensure there is only one parent session
+        String[] sessionIds = getDevice().executeShellCommand(
+                "pm get-stagedsessions --only-ready --only-parent --only-sessionid").split("\n");
+        assertThat(sessionIds.length).isEqualTo(1);
+        // Ensure there are two children session
+        sessionIds = getDevice().executeShellCommand(
+                "pm get-stagedsessions --only-ready --only-sessionid").split("\n");
+        assertThat(sessionIds.length).isEqualTo(3);
+    }
+
     private void restartSystemServer() throws Exception {
         // Restart the system server
         long oldStartTime = getDevice().getProcessByName("system_server").getStartTime();