Added UserManagerInternal.getUserAssignedToDisplay(displayId).

Test: atest FrameworksMockingServicesTests:com.android.server.pm.UserManagerServiceTest UserManagerInternalTest

Bug: 244331974
Change-Id: I06b933f15bf6903c40839ed1090439f2a99da5a6

5 files changed, 128 insertions, 2 deletions

diff --git a/services/core/java/com/android/server/pm/UserManagerInternal.java b/services/core/java/com/android/server/pm/UserManagerInternal.java
index 62655843a5f2..b9a119504d1f 100644
--- a/services/core/java/com/android/server/pm/UserManagerInternal.java
+++ b/services/core/java/com/android/server/pm/UserManagerInternal.java
@@ -352,10 +352,22 @@ public abstract class UserManagerInternal {
      * Returns the display id assigned to the user, or {@code Display.INVALID_DISPLAY} if the
      * user is not assigned to any display.
      *
-     * <p>The current foreground user is associated with the main display, while other users would
-     * only assigned to a display if they were started with
+     * <p>The current foreground user is associated with the
+     * {@link android.view.Display#DEFAULT_DISPLAY default display}, while other users would only be
+     * assigned to a display if they were started with
      * {@code ActivityManager.startUserInBackgroundOnSecondaryDisplay()}. If the user is a profile
      * and is running, it's assigned to its parent display.
      */
     public abstract int getDisplayAssignedToUser(@UserIdInt int userId);
+
+    /**
+     * Returns the main user (i.e., not a profile) that is assigned to the display, or the
+     * {@link android.app.ActivityManager#getCurrentUser() current foreground user} if no user is
+     * associated with the display.
+     *
+     * <p>The {@link android.view.Display#DEFAULT_DISPLAY default display} is always assigned to
+     * the current foreground user, while other displays would be associated with the user that was
+     * started with {@code ActivityManager.startUserInBackgroundOnSecondaryDisplay()}.
+     */
+    public abstract @UserIdInt int getUserAssignedToDisplay(int displayId);
 }
diff --git a/services/core/java/com/android/server/pm/UserManagerService.java b/services/core/java/com/android/server/pm/UserManagerService.java
index 1e30757b8417..964b367c010f 100644
--- a/services/core/java/com/android/server/pm/UserManagerService.java
+++ b/services/core/java/com/android/server/pm/UserManagerService.java
@@ -1778,6 +1778,43 @@ public class UserManagerService extends IUserManager.Stub {
         }
     }
 
+    @VisibleForTesting
+    int getUserAssignedToDisplay(int displayId) {
+        if (displayId == Display.DEFAULT_DISPLAY) {
+            return getCurrentUserId();
+        }
+
+        if (!mUsersOnSecondaryDisplaysEnabled) {
+            int currentUserId = getCurrentUserId();
+            Slogf.w(LOG_TAG, "getUsersAssignedToDisplay(%d) called with non-DEFAULT_DISPLAY on "
+                    + "system that doesn't support that; returning current user (%d)", displayId,
+                    currentUserId);
+            return currentUserId;
+        }
+
+        synchronized (mUsersOnSecondaryDisplays) {
+            for (int i = 0; i < mUsersOnSecondaryDisplays.size(); i++) {
+                if (mUsersOnSecondaryDisplays.valueAt(i) != displayId) {
+                    continue;
+                }
+                int userId = mUsersOnSecondaryDisplays.keyAt(i);
+                if (!isProfileUnchecked(userId)) {
+                    return userId;
+                } else if (DBG_MUMD) {
+                    Slogf.d(LOG_TAG, "getUserAssignedToDisplay(%d): skipping user %d because it's "
+                            + "a profile", displayId, userId);
+                }
+            }
+        }
+
+        int currentUserId = getCurrentUserId();
+        if (DBG_MUMD) {
+            Slogf.d(LOG_TAG, "getUserAssignedToDisplay(%d): no user assigned to display, returning "
+                    + "current user (%d) instead", displayId, currentUserId);
+        }
+        return currentUserId;
+    }
+
     /**
      * Gets the current user id, calling {@link ActivityManagerInternal} directly (and without
      * performing any permission check).
@@ -6827,6 +6864,11 @@ public class UserManagerService extends IUserManager.Stub {
         public int getDisplayAssignedToUser(int userId) {
             return UserManagerService.this.getDisplayAssignedToUser(userId);
         }
+
+        @Override
+        public int getUserAssignedToDisplay(int displayId) {
+            return UserManagerService.this.getUserAssignedToDisplay(displayId);
+        }
     } // class LocalService
 
     /**
diff --git a/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerInternalTest.java b/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerInternalTest.java
index 1cff685522a2..86a5c90596ff 100644
--- a/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerInternalTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerInternalTest.java
@@ -229,4 +229,9 @@ public final class UserManagerInternalTest extends UserManagerServiceOrInternalT
     protected int getDisplayAssignedToUser(int userId) {
         return mUmi.getDisplayAssignedToUser(userId);
     }
+
+    @Override
+    protected int getUserAssignedToDisplay(int displayId) {
+        return mUmi.getUserAssignedToDisplay(displayId);
+    }
 }
diff --git a/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerServiceOrInternalTestCase.java b/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerServiceOrInternalTestCase.java
index 4f61b8fff995..991053a120ef 100644
--- a/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerServiceOrInternalTestCase.java
+++ b/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerServiceOrInternalTestCase.java
@@ -194,6 +194,7 @@ abstract class UserManagerServiceOrInternalTestCase extends ExtendedMockitoTestC
     protected abstract boolean isUserVisible(int userId);
     protected abstract boolean isUserVisibleOnDisplay(int userId, int displayId);
     protected abstract int getDisplayAssignedToUser(int userId);
+    protected abstract int getUserAssignedToDisplay(int displayId);
 
     /////////////////////////////////
     // Tests for the above methods //
@@ -427,6 +428,67 @@ abstract class UserManagerServiceOrInternalTestCase extends ExtendedMockitoTestC
     // NOTE: we don't need to add tests for profiles (started / stopped profiles of bg user), as
     // getDisplayAssignedToUser() for bg users relies only on the user / display assignments
 
+    @Test
+    public void testGetUserAssignedToDisplay_invalidDisplay() {
+        mockCurrentUser(USER_ID);
+
+        assertWithMessage("getUserAssignedToDisplay(%s)", INVALID_DISPLAY)
+                .that(getUserAssignedToDisplay(INVALID_DISPLAY)).isEqualTo(USER_ID);
+    }
+
+    @Test
+    public void testGetUserAssignedToDisplay_defaultDisplay() {
+        mockCurrentUser(USER_ID);
+
+        assertWithMessage("getUserAssignedToDisplay(%s)", DEFAULT_DISPLAY)
+                .that(getUserAssignedToDisplay(DEFAULT_DISPLAY)).isEqualTo(USER_ID);
+    }
+
+    @Test
+    public void testGetUserAssignedToDisplay_secondaryDisplay() {
+        mockCurrentUser(USER_ID);
+
+        assertWithMessage("getUserAssignedToDisplay(%s)", SECONDARY_DISPLAY_ID)
+                .that(getUserAssignedToDisplay(SECONDARY_DISPLAY_ID)).isEqualTo(USER_ID);
+    }
+
+    @Test
+    public void testGetUserAssignedToDisplay_mumd_bgUserOnSecondaryDisplay() {
+        enableUsersOnSecondaryDisplays();
+        mockCurrentUser(OTHER_USER_ID);
+        assignUserToDisplay(USER_ID, SECONDARY_DISPLAY_ID);
+
+        assertWithMessage("getUserAssignedToDisplay(%s)", SECONDARY_DISPLAY_ID)
+                .that(getUserAssignedToDisplay(SECONDARY_DISPLAY_ID)).isEqualTo(USER_ID);
+    }
+
+    @Test
+    public void testGetUserAssignedToDisplay_mumd_noUserOnSecondaryDisplay() {
+        enableUsersOnSecondaryDisplays();
+        mockCurrentUser(USER_ID);
+
+        assertWithMessage("getUserAssignedToDisplay(%s)", SECONDARY_DISPLAY_ID)
+                .that(getUserAssignedToDisplay(SECONDARY_DISPLAY_ID)).isEqualTo(USER_ID);
+    }
+
+    // TODO(b/244644281): scenario below shouldn't happen on "real life", as the profile cannot be
+    // started on secondary display if its parent isn't, so we might need to remove (or refactor
+    // this test) if/when the underlying logic changes
+    @Test
+    public void testGetUserAssignedToDisplay_mumd_profileOnSecondaryDisplay() {
+        enableUsersOnSecondaryDisplays();
+        addDefaultProfileAndParent();
+        mockCurrentUser(USER_ID);
+        assignUserToDisplay(PROFILE_USER_ID, SECONDARY_DISPLAY_ID);
+
+        assertWithMessage("getUserAssignedToDisplay(%s)", SECONDARY_DISPLAY_ID)
+                .that(getUserAssignedToDisplay(SECONDARY_DISPLAY_ID)).isEqualTo(USER_ID);
+    }
+
+    // NOTE: we don't need to add tests for profiles (started / stopped profiles of bg user), as
+    // getUserAssignedToDisplay() for bg users relies only on the user / display assignments
+
+
     ///////////////////////////////////////////
     // Helper methods exposed to sub-classes //
     ///////////////////////////////////////////
diff --git a/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerServiceTest.java b/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerServiceTest.java
index b335a34f6f14..8388a70a87e5 100644
--- a/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerServiceTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/pm/UserManagerServiceTest.java
@@ -102,4 +102,9 @@ public final class UserManagerServiceTest extends UserManagerServiceOrInternalTe
     protected int getDisplayAssignedToUser(int userId) {
         return mUms.getDisplayAssignedToUser(userId);
     }
+
+    @Override
+    protected int getUserAssignedToDisplay(int displayId) {
+        return mUms.getUserAssignedToDisplay(displayId);
+    }
 }