From e2043617199c6044f487f2a78eb658fddbc91768 Mon Sep 17 00:00:00 2001
From: jovanak <jovanak@google.com>
Date: Thu, 31 Oct 2019 16:49:24 -0700
Subject: [PATCH] Add wait flag to am switch-user command.

Previously, command returns right away. CTS tests that use it are forced to
loop waiting for the change in am get-current-user to see if the switch
completed. This method is unreliable because it returns the new user id before
UserController.mCurrentUser updated.

Adding "-w" to the switch command will make it wait for the switch to be
completed.

Fixes:141998718
Test: manual switch-user tests with -w flag
Change-Id: I41bbcff0c79677ba3bf6b58e1fc96cd13a681b33
---
 .../am/ActivityManagerShellCommand.java       | 44 ++++++++++++++++++-
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/services/core/java/com/android/server/am/ActivityManagerShellCommand.java b/services/core/java/com/android/server/am/ActivityManagerShellCommand.java
index 59acdcf4a8757..15a6468e137a2 100644
--- a/services/core/java/com/android/server/am/ActivityManagerShellCommand.java
+++ b/services/core/java/com/android/server/am/ActivityManagerShellCommand.java
@@ -37,6 +37,7 @@ import android.app.IStopUserCallback;
 import android.app.IUidObserver;
 import android.app.KeyguardManager;
 import android.app.ProfilerInfo;
+import android.app.UserSwitchObserver;
 import android.app.WaitResult;
 import android.app.usage.AppStandbyInfo;
 import android.app.usage.ConfigurationStats;
@@ -1729,6 +1730,30 @@ final class ActivityManagerShellCommand extends ShellCommand {
         return 0;
     }
 
+    private void switchUserAndWaitForComplete(int userId) throws RemoteException {
+        // Register switch observer.
+        final CountDownLatch switchLatch = new CountDownLatch(1);
+        mInterface.registerUserSwitchObserver(
+                new UserSwitchObserver() {
+                    @Override
+                    public void onUserSwitchComplete(int newUserId) {
+                        if (userId == newUserId) {
+                            switchLatch.countDown();
+                        }
+                    }
+                }, ActivityManagerShellCommand.class.getName());
+
+        // Switch.
+        mInterface.switchUser(userId);
+
+        // Wait.
+        try {
+            switchLatch.await(USER_OPERATION_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+            getErrPrintWriter().println("Thread interrupted unexpectedly.");
+        }
+    }
+
     int runSwitchUser(PrintWriter pw) throws RemoteException {
         UserManager userManager = mInternal.mContext.getSystemService(UserManager.class);
         final int userSwitchable = userManager.getUserSwitchability();
@@ -1736,8 +1761,23 @@ final class ActivityManagerShellCommand extends ShellCommand {
             getErrPrintWriter().println("Error: " + userSwitchable);
             return -1;
         }
-        String user = getNextArgRequired();
-        mInterface.switchUser(Integer.parseInt(user));
+        boolean wait = false;
+        String opt;
+        while ((opt = getNextOption()) != null) {
+            if ("-w".equals(opt)) {
+                wait = true;
+            } else {
+                getErrPrintWriter().println("Error: unknown option: " + opt);
+                return -1;
+            }
+        }
+
+        int userId = Integer.parseInt(getNextArgRequired());
+        if (wait) {
+            switchUserAndWaitForComplete(userId);
+        } else {
+            mInterface.switchUser(userId);
+        }
         return 0;
     }