1 /*
2  * Copyright (C) 2020 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.documentsui;
18 
19 import static com.android.documentsui.DevicePolicyResources.Drawables.Style.SOLID_COLORED;
20 import static com.android.documentsui.DevicePolicyResources.Drawables.WORK_PROFILE_ICON;
21 import static com.android.documentsui.DevicePolicyResources.Strings.PERSONAL_TAB;
22 import static com.android.documentsui.DevicePolicyResources.Strings.WORK_TAB;
23 
24 import static com.google.common.truth.Truth.assertWithMessage;
25 
26 import static org.mockito.ArgumentMatchers.any;
27 import static org.mockito.ArgumentMatchers.eq;
28 import static org.mockito.Mockito.mock;
29 import static org.mockito.Mockito.when;
30 
31 import android.Manifest;
32 import android.app.admin.DevicePolicyManager;
33 import android.app.admin.DevicePolicyResourcesManager;
34 import android.content.Context;
35 import android.content.Intent;
36 import android.content.pm.PackageManager;
37 import android.content.pm.ResolveInfo;
38 import android.content.pm.UserProperties;
39 import android.content.res.Resources;
40 import android.graphics.drawable.Drawable;
41 import android.os.UserHandle;
42 import android.os.UserManager;
43 
44 import androidx.test.filters.SdkSuppress;
45 import androidx.test.filters.SmallTest;
46 import androidx.test.platform.app.InstrumentationRegistry;
47 
48 import com.android.documentsui.base.UserId;
49 import com.android.documentsui.testing.UserManagers;
50 import com.android.documentsui.util.VersionUtils;
51 import com.android.modules.utils.build.SdkLevel;
52 
53 import com.google.common.collect.Lists;
54 
55 import org.junit.Before;
56 import org.junit.Test;
57 
58 import java.util.ArrayList;
59 import java.util.HashMap;
60 import java.util.List;
61 import java.util.Map;
62 
63 @SmallTest
64 @SdkSuppress(minSdkVersion = 31, codeName = "S")
65 public class UserManagerStateTest {
66 
67     private static final String PERSONAL = "Personal";
68     private static final String WORK = "Work";
69     private static final String PRIVATE = "Private";
70 
71     private final UserHandle mSystemUser = UserHandle.SYSTEM;
72     private final UserHandle mManagedUser = UserHandle.of(100);
73     private final UserHandle mPrivateUser = UserHandle.of(101);
74     private final UserHandle mOtherUser = UserHandle.of(102);
75     private final UserHandle mNormalUser = UserHandle.of(103);
76 
77     private final ResolveInfo mMockInfo1 = mock(ResolveInfo.class);
78     private final ResolveInfo mMockInfo2 = mock(ResolveInfo.class);
79     private final ResolveInfo mMockInfo3 = mock(ResolveInfo.class);
80 
81     private final Context mMockContext = mock(Context.class);
82     private final Intent mMockIntent = mock(Intent.class);
83     private final UserManager mMockUserManager = UserManagers.create();
84     private final PackageManager mMockPackageManager = mock(PackageManager.class);
85     private final DevicePolicyManager mDevicePolicyManager = mock(DevicePolicyManager.class);
86     private UserManagerState mUserManagerState;
87 
88     @Before
setup()89     public void setup() throws Exception {
90         when(mMockContext.getApplicationContext()).thenReturn(mMockContext);
91 
92         when(mMockUserManager.isManagedProfile(mManagedUser.getIdentifier())).thenReturn(true);
93         when(mMockUserManager.isManagedProfile(mSystemUser.getIdentifier())).thenReturn(false);
94         when(mMockUserManager.isManagedProfile(mPrivateUser.getIdentifier())).thenReturn(false);
95         when(mMockUserManager.isManagedProfile(mOtherUser.getIdentifier())).thenReturn(false);
96 
97         if (SdkLevel.isAtLeastV()) {
98             UserProperties systemUserProperties = new UserProperties.Builder()
99                     .setShowInSharingSurfaces(UserProperties.SHOW_IN_SHARING_SURFACES_SEPARATE)
100                     .setCrossProfileContentSharingStrategy(
101                             UserProperties.CROSS_PROFILE_CONTENT_SHARING_NO_DELEGATION)
102                     .build();
103             UserProperties managedUserProperties = new UserProperties.Builder()
104                     .setShowInSharingSurfaces(UserProperties.SHOW_IN_SHARING_SURFACES_SEPARATE)
105                     .setCrossProfileContentSharingStrategy(
106                             UserProperties.CROSS_PROFILE_CONTENT_SHARING_NO_DELEGATION)
107                     .setShowInQuietMode(UserProperties.SHOW_IN_QUIET_MODE_PAUSED)
108                     .build();
109             UserProperties privateUserProperties = new UserProperties.Builder()
110                     .setShowInSharingSurfaces(UserProperties.SHOW_IN_SHARING_SURFACES_SEPARATE)
111                     .setCrossProfileContentSharingStrategy(
112                             UserProperties.CROSS_PROFILE_CONTENT_SHARING_DELEGATE_FROM_PARENT)
113                     .setShowInQuietMode(UserProperties.SHOW_IN_QUIET_MODE_HIDDEN)
114                     .build();
115             UserProperties otherUserProperties = new UserProperties.Builder()
116                     .setShowInSharingSurfaces(UserProperties.SHOW_IN_SHARING_SURFACES_WITH_PARENT)
117                     .setCrossProfileContentSharingStrategy(
118                             UserProperties.CROSS_PROFILE_CONTENT_SHARING_DELEGATE_FROM_PARENT)
119                     .build();
120             UserProperties normalUserProperties = new UserProperties.Builder()
121                     .setShowInSharingSurfaces(UserProperties.SHOW_IN_SHARING_SURFACES_NO)
122                     .setCrossProfileContentSharingStrategy(
123                             UserProperties.CROSS_PROFILE_CONTENT_SHARING_DELEGATE_FROM_PARENT)
124                     .build();
125             when(mMockUserManager.getUserProperties(mSystemUser)).thenReturn(systemUserProperties);
126             when(mMockUserManager.getUserProperties(mManagedUser)).thenReturn(
127                     managedUserProperties);
128             when(mMockUserManager.getUserProperties(mPrivateUser)).thenReturn(
129                     privateUserProperties);
130             when(mMockUserManager.getUserProperties(mOtherUser)).thenReturn(otherUserProperties);
131             when(mMockUserManager.getUserProperties(mNormalUser)).thenReturn(normalUserProperties);
132         }
133 
134         when(mMockUserManager.getProfileParent(mSystemUser)).thenReturn(null);
135         when(mMockUserManager.getProfileParent(mManagedUser)).thenReturn(mSystemUser);
136         when(mMockUserManager.getProfileParent(mPrivateUser)).thenReturn(mSystemUser);
137         when(mMockUserManager.getProfileParent(mOtherUser)).thenReturn(mSystemUser);
138         when(mMockUserManager.getProfileParent(mNormalUser)).thenReturn(null);
139 
140         if (SdkLevel.isAtLeastR()) {
141             when(mMockInfo1.isCrossProfileIntentForwarderActivity()).thenReturn(true);
142             when(mMockInfo2.isCrossProfileIntentForwarderActivity()).thenReturn(false);
143             when(mMockInfo3.isCrossProfileIntentForwarderActivity()).thenReturn(false);
144         }
145 
146         when(mMockContext.getPackageManager()).thenReturn(mMockPackageManager);
147         when(mMockContext.getSystemServiceName(UserManager.class)).thenReturn("mMockUserManager");
148         when(mMockContext.getSystemService(UserManager.class)).thenReturn(mMockUserManager);
149         when(mMockContext.getSystemServiceName(DevicePolicyManager.class))
150                 .thenReturn(Context.DEVICE_POLICY_SERVICE);
151         when(mMockContext.getSystemService(Context.DEVICE_POLICY_SERVICE))
152                 .thenReturn(mDevicePolicyManager);
153         when(mMockContext.getResources()).thenReturn(
154                 InstrumentationRegistry.getInstrumentation().getTargetContext().getResources());
155     }
156 
157     @Test
testGetUserIds_onlySystemUser_returnsSystemUser()158     public void testGetUserIds_onlySystemUser_returnsSystemUser() {
159         UserId currentUser = UserId.of(mSystemUser);
160         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser));
161 
162         assertWithMessage("getUserIds returns unexpected list of user ids")
163                 .that(mUserManagerState.getUserIds()).containsExactly(UserId.of(mSystemUser));
164     }
165 
166     @Test
testGetUserIds_allProfilesCurrentUserSystem_allShowInSharingSurfacesSeparate()167     public void testGetUserIds_allProfilesCurrentUserSystem_allShowInSharingSurfacesSeparate() {
168         if (!SdkLevel.isAtLeastV()) return;
169         UserId currentUser = UserId.of(mSystemUser);
170         initializeUserManagerState(currentUser,
171                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser, mOtherUser,
172                         mNormalUser));
173 
174         assertWithMessage("getUserIds returns unexpected list of user ids")
175                 .that(mUserManagerState.getUserIds())
176                 .containsExactly(UserId.of(mSystemUser), UserId.of(mManagedUser),
177                         UserId.of(mPrivateUser));
178     }
179 
180     @Test
testGetUserIds_allProfilesCurrentUserManaged_allShowInSharingSurfacesSeparate()181     public void testGetUserIds_allProfilesCurrentUserManaged_allShowInSharingSurfacesSeparate() {
182         if (!SdkLevel.isAtLeastV()) return;
183         UserId currentUser = UserId.of(mManagedUser);
184         initializeUserManagerState(currentUser,
185                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser, mOtherUser,
186                         mNormalUser));
187 
188         assertWithMessage("getUserIds returns unexpected list of user ids")
189                 .that(mUserManagerState.getUserIds())
190                 .containsExactly(UserId.of(mSystemUser), UserId.of(mManagedUser),
191                         UserId.of(mPrivateUser));
192     }
193 
194     @Test
testGetUserIds_allProfilesCurrentUserPrivate_allShowInSharingSurfacesSeparate()195     public void testGetUserIds_allProfilesCurrentUserPrivate_allShowInSharingSurfacesSeparate() {
196         if (!SdkLevel.isAtLeastV()) return;
197         UserId currentUser = UserId.of(mPrivateUser);
198         initializeUserManagerState(currentUser,
199                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser, mOtherUser));
200 
201         assertWithMessage("getUserIds returns unexpected list of user ids")
202                 .that(mUserManagerState.getUserIds())
203                 .containsExactly(UserId.of(mSystemUser), UserId.of(mManagedUser),
204                         UserId.of(mPrivateUser));
205     }
206 
207     @Test
testGetUserIds_systemAndManagedUserCurrentUserSystem_returnsBoth()208     public void testGetUserIds_systemAndManagedUserCurrentUserSystem_returnsBoth() {
209         UserId currentUser = UserId.of(mSystemUser);
210         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mManagedUser));
211 
212         assertWithMessage("getUserIds returns unexpected list of user ids")
213                 .that(mUserManagerState.getUserIds())
214                 .containsExactly(UserId.of(mSystemUser), UserId.of(mManagedUser));
215     }
216 
217     @Test
testGetUserIds_systemAndManagedUserCurrentUserManaged_returnsBoth()218     public void testGetUserIds_systemAndManagedUserCurrentUserManaged_returnsBoth() {
219         UserId currentUser = UserId.of(mManagedUser);
220         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mManagedUser));
221 
222         assertWithMessage("getUserIds returns unexpected list of user ids")
223                 .that(mUserManagerState.getUserIds())
224                 .containsExactly(UserId.of(mSystemUser), UserId.of(mManagedUser));
225     }
226 
227     @Test
testGetUserIds_systemAndPrivateUserCurrentUserSystem_returnsBoth()228     public void testGetUserIds_systemAndPrivateUserCurrentUserSystem_returnsBoth() {
229         if (!SdkLevel.isAtLeastV()) return;
230         UserId currentUser = UserId.of(mSystemUser);
231         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mPrivateUser));
232 
233         assertWithMessage("getUserIds returns unexpected list of user ids")
234                 .that(mUserManagerState.getUserIds())
235                 .containsExactly(UserId.of(mSystemUser), UserId.of(mPrivateUser));
236     }
237 
238     @Test
testGetUserIds_systemAndPrivateUserCurrentUserPrivate_returnsBoth()239     public void testGetUserIds_systemAndPrivateUserCurrentUserPrivate_returnsBoth() {
240         if (!SdkLevel.isAtLeastV()) return;
241         UserId currentUser = UserId.of(mPrivateUser);
242         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mPrivateUser));
243 
244         assertWithMessage("getUserIds returns unexpected list of user ids")
245                 .that(mUserManagerState.getUserIds())
246                 .containsExactly(UserId.of(mSystemUser), UserId.of(mPrivateUser));
247     }
248 
249     @Test
testGetUserIds_systemAndOtherUserCurrentUserOtherPreV_returnsCurrentUser()250     public void testGetUserIds_systemAndOtherUserCurrentUserOtherPreV_returnsCurrentUser() {
251         if (SdkLevel.isAtLeastV()) return;
252         UserId currentUser = UserId.of(mOtherUser);
253         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mOtherUser));
254 
255         assertWithMessage("getUserIds returns unexpected list of user ids")
256                 .that(mUserManagerState.getUserIds())
257                 .containsExactly(currentUser);
258     }
259 
260     @Test
testGetUserIds_systemAndOtherUserCurrentUserOtherPostV_returnsSystemUser()261     public void testGetUserIds_systemAndOtherUserCurrentUserOtherPostV_returnsSystemUser() {
262         if (!SdkLevel.isAtLeastV()) return;
263         UserId currentUser = UserId.of(mOtherUser);
264         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mOtherUser));
265 
266         assertWithMessage("getUserIds returns unexpected list of user ids")
267                 .that(mUserManagerState.getUserIds())
268                 .containsExactly(UserId.of(mSystemUser));
269     }
270 
271     @Test
testGetUserIds_normalAndOtherUserCurrentUserNormal_returnsCurrentUser()272     public void testGetUserIds_normalAndOtherUserCurrentUserNormal_returnsCurrentUser() {
273         // since both users do not have show in sharing surfaces separate, returns current user
274         UserId currentUser = UserId.of(mNormalUser);
275         initializeUserManagerState(currentUser, Lists.newArrayList(mOtherUser, mNormalUser));
276 
277         assertWithMessage("getUserIds returns unexpected list of user ids")
278                 .that(mUserManagerState.getUserIds())
279                 .containsExactly(UserId.of(mNormalUser));
280     }
281 
282     @Test
testGetUserIds_systemAndManagedUserCurrentUserSystem_returnsBothInOrder()283     public void testGetUserIds_systemAndManagedUserCurrentUserSystem_returnsBothInOrder() {
284         // Returns the both if there are system and managed users.
285         if (SdkLevel.isAtLeastV()) return;
286         UserId currentUser = UserId.of(mSystemUser);
287         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mManagedUser));
288         assertWithMessage("getUserIds returns unexpected list of user ids")
289                 .that(mUserManagerState.getUserIds())
290                 .containsExactly(UserId.of(mSystemUser), UserId.of(mManagedUser)).inOrder();
291     }
292 
293     @Test
testGetUserIds_systemAndManagedUserCurrentUserManaged_returnsBothInOrder()294     public void testGetUserIds_systemAndManagedUserCurrentUserManaged_returnsBothInOrder() {
295         // Returns the both if there are system and managed users.
296         if (SdkLevel.isAtLeastV()) return;
297         UserId currentUser = UserId.of(mManagedUser);
298         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mManagedUser));
299         assertWithMessage("getUserIds returns unexpected list of user ids")
300                 .that(mUserManagerState.getUserIds())
301                 .containsExactly(UserId.of(mSystemUser), UserId.of(mManagedUser)).inOrder();
302     }
303 
304     @Test
testGetUserIds_managedAndSystemUserCurrentUserSystem_returnsBothInOrder()305     public void testGetUserIds_managedAndSystemUserCurrentUserSystem_returnsBothInOrder() {
306         // Returns the both if there are system and managed users, regardless of input list order.
307         if (SdkLevel.isAtLeastV()) return;
308         UserId currentUser = UserId.of(mSystemUser);
309         initializeUserManagerState(currentUser, Lists.newArrayList(mManagedUser, mSystemUser));
310         assertWithMessage("getUserIds returns unexpected list of user ids")
311                 .that(mUserManagerState.getUserIds())
312                 .containsExactly(UserId.of(mSystemUser), UserId.of(mManagedUser)).inOrder();
313     }
314 
315     @Test
testGetUserIds_otherAndManagedUserCurrentUserOtherPreV_returnsCurrentUser()316     public void testGetUserIds_otherAndManagedUserCurrentUserOtherPreV_returnsCurrentUser() {
317         // When there is no system user, returns the current user.
318         // This is a case theoretically can happen but we don't expect. So we return the current
319         // user only.
320         if (SdkLevel.isAtLeastV()) return;
321         UserId currentUser = UserId.of(mOtherUser);
322         initializeUserManagerState(currentUser, Lists.newArrayList(mOtherUser, mManagedUser));
323         assertWithMessage("getUserIds returns unexpected list of user ids")
324                 .that(mUserManagerState.getUserIds()).containsExactly(currentUser);
325     }
326 
327     @Test
testGetUserIds_otherAndManagedUserCurrentUserOtherPostV_returnsManagedUser()328     public void testGetUserIds_otherAndManagedUserCurrentUserOtherPostV_returnsManagedUser() {
329         // Only the users with show in sharing surfaces separate are eligible to be returned
330         if (!SdkLevel.isAtLeastV()) return;
331         UserId currentUser = UserId.of(mOtherUser);
332         initializeUserManagerState(currentUser, Lists.newArrayList(mOtherUser, mManagedUser));
333         assertWithMessage("getUserIds returns unexpected list of user ids")
334                 .that(mUserManagerState.getUserIds()).containsExactly(UserId.of(mManagedUser));
335     }
336 
337     @Test
testGetUserIds_otherAndManagedUserCurrentUserManaged_returnsCurrentUser()338     public void testGetUserIds_otherAndManagedUserCurrentUserManaged_returnsCurrentUser() {
339         // When there is no system user, returns the current user.
340         // This is a case theoretically can happen, but we don't expect. So we return the current
341         // user only.
342         UserId currentUser = UserId.of(mManagedUser);
343         initializeUserManagerState(currentUser, Lists.newArrayList(mOtherUser, mManagedUser));
344         assertWithMessage("getUserIds returns unexpected list of user ids")
345                 .that(mUserManagerState.getUserIds()).containsExactly(currentUser);
346     }
347 
348     @Test
testGetUserIds_unsupportedDeviceCurrent_returnsCurrentUser()349     public void testGetUserIds_unsupportedDeviceCurrent_returnsCurrentUser() {
350         // This test only tests for Android R or later. This test case always passes before R.
351         if (VersionUtils.isAtLeastR()) {
352             // When permission is denied, only returns the current user.
353             when(mMockContext.checkSelfPermission(Manifest.permission.INTERACT_ACROSS_USERS))
354                     .thenReturn(PackageManager.PERMISSION_DENIED);
355             UserId currentUser = UserId.of(mSystemUser);
356             when(mMockUserManager.getUserProfiles()).thenReturn(
357                     Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser, mOtherUser));
358             mUserManagerState = UserManagerState.create(mMockContext);
359             assertWithMessage("Unsupported device should have returned only the current user")
360                     .that(mUserManagerState.getUserIds()).containsExactly(currentUser);
361         }
362     }
363 
364     @Test
testGetUserIds_returnCachedList()365     public void testGetUserIds_returnCachedList() {
366         // Returns all three if there are system, managed and private users.
367         UserId currentUser = UserId.of(mSystemUser);
368         initializeUserManagerState(currentUser,
369                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser, mOtherUser));
370         assertWithMessage("getUserIds does not return cached instance")
371                 .that(mUserManagerState.getUserIds())
372                 .isSameInstanceAs(mUserManagerState.getUserIds());
373     }
374 
375     @Test
testGetCanForwardToProfileIdMap_systemUserCanForwardToAll()376     public void testGetCanForwardToProfileIdMap_systemUserCanForwardToAll() {
377         UserId currentUser = UserId.of(mSystemUser);
378         final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo1, mMockInfo2);
379         if (SdkLevel.isAtLeastV()) {
380             initializeUserManagerState(currentUser,
381                     Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
382             when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
383                     PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
384                     mMockResolveInfoList);
385         } else {
386             initializeUserManagerState(currentUser,
387                     Lists.newArrayList(mSystemUser, mManagedUser));
388             when(mMockPackageManager.queryIntentActivities(mMockIntent,
389                     PackageManager.MATCH_DEFAULT_ONLY)).thenReturn(mMockResolveInfoList);
390         }
391 
392         Map<UserId, Boolean> expectedCanForwardToProfileIdMap = new HashMap<>();
393         expectedCanForwardToProfileIdMap.put(UserId.of(mSystemUser), true);
394         expectedCanForwardToProfileIdMap.put(UserId.of(mManagedUser), true);
395         if (SdkLevel.isAtLeastV()) {
396             expectedCanForwardToProfileIdMap.put(UserId.of(mPrivateUser), true);
397         }
398 
399         assertWithMessage("getCanForwardToProfileIdMap returns incorrect mappings")
400                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent))
401                 .isEqualTo(expectedCanForwardToProfileIdMap);
402     }
403 
404     @Test
testGetCanForwardToProfileIdMap_systemUserCanForwardToManaged()405     public void testGetCanForwardToProfileIdMap_systemUserCanForwardToManaged() {
406         UserId currentUser = UserId.of(mSystemUser);
407         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mManagedUser));
408         final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo1, mMockInfo2);
409         if (SdkLevel.isAtLeastV()) {
410             when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
411                     PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
412                     mMockResolveInfoList);
413         } else {
414             when(mMockPackageManager.queryIntentActivities(mMockIntent,
415                     PackageManager.MATCH_DEFAULT_ONLY)).thenReturn(mMockResolveInfoList);
416         }
417 
418         Map<UserId, Boolean> expectedCanForwardToProfileIdMap = new HashMap<>();
419         expectedCanForwardToProfileIdMap.put(UserId.of(mSystemUser), true);
420         expectedCanForwardToProfileIdMap.put(UserId.of(mManagedUser), true);
421 
422         assertWithMessage("getCanForwardToProfileIdMap returns incorrect mappings")
423                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent))
424                 .isEqualTo(expectedCanForwardToProfileIdMap);
425     }
426 
427     @Test
testGetCanForwardToProfileIdMap_systemUserCanAlwaysForwardToPrivate()428     public void testGetCanForwardToProfileIdMap_systemUserCanAlwaysForwardToPrivate() {
429         if (!SdkLevel.isAtLeastV()) return;
430         UserId currentUser = UserId.of(mSystemUser);
431         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mPrivateUser));
432 
433         Map<UserId, Boolean> expectedCanForwardToProfileIdMap = new HashMap<>();
434         expectedCanForwardToProfileIdMap.put(UserId.of(mSystemUser), true);
435         expectedCanForwardToProfileIdMap.put(UserId.of(mPrivateUser), true);
436 
437         assertWithMessage("getCanForwardToProfileIdMap returns incorrect mappings")
438                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent))
439                 .isEqualTo(expectedCanForwardToProfileIdMap);
440     }
441 
442     @Test
testGetCanForwardToProfileIdMap_systemUserCanNotForwardToManagedUser()443     public void testGetCanForwardToProfileIdMap_systemUserCanNotForwardToManagedUser() {
444         UserId currentUser = UserId.of(mSystemUser);
445         final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo2, mMockInfo3);
446         if (SdkLevel.isAtLeastV()) {
447             initializeUserManagerState(currentUser,
448                     Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
449             when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
450                     PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
451                     mMockResolveInfoList);
452         } else {
453             initializeUserManagerState(currentUser,
454                     Lists.newArrayList(mSystemUser, mManagedUser));
455             when(mMockPackageManager.queryIntentActivities(mMockIntent,
456                     PackageManager.MATCH_DEFAULT_ONLY)).thenReturn(mMockResolveInfoList);
457         }
458 
459         Map<UserId, Boolean> expectedCanForwardToProfileIdMap = new HashMap<>();
460         expectedCanForwardToProfileIdMap.put(UserId.of(mSystemUser), true);
461         expectedCanForwardToProfileIdMap.put(UserId.of(mManagedUser), false);
462         if (SdkLevel.isAtLeastV()) {
463             expectedCanForwardToProfileIdMap.put(UserId.of(mPrivateUser), true);
464         }
465 
466         assertWithMessage("getCanForwardToProfileIdMap returns incorrect mappings")
467                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent))
468                 .isEqualTo(expectedCanForwardToProfileIdMap);
469     }
470 
471     @Test
testGetCanForwardToProfileIdMap_managedCanForwardToAll()472     public void testGetCanForwardToProfileIdMap_managedCanForwardToAll() {
473         UserId currentUser = UserId.of(mManagedUser);
474         final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo1, mMockInfo2);
475         if (SdkLevel.isAtLeastV()) {
476             initializeUserManagerState(currentUser,
477                     Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
478             when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
479                     PackageManager.MATCH_DEFAULT_ONLY, mManagedUser)).thenReturn(
480                     mMockResolveInfoList);
481         } else {
482             initializeUserManagerState(currentUser,
483                     Lists.newArrayList(mSystemUser, mManagedUser));
484             when(mMockPackageManager.queryIntentActivities(mMockIntent,
485                     PackageManager.MATCH_DEFAULT_ONLY)).thenReturn(mMockResolveInfoList);
486         }
487 
488         Map<UserId, Boolean> expectedCanForwardToProfileIdMap = new HashMap<>();
489         expectedCanForwardToProfileIdMap.put(UserId.of(mSystemUser), true);
490         expectedCanForwardToProfileIdMap.put(UserId.of(mManagedUser), true);
491         if (SdkLevel.isAtLeastV()) {
492             expectedCanForwardToProfileIdMap.put(UserId.of(mPrivateUser), true);
493         }
494 
495         assertWithMessage("getCanForwardToProfileIdMap returns incorrect mappings")
496                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent))
497                 .isEqualTo(expectedCanForwardToProfileIdMap);
498     }
499 
500     @Test
testGetCanForwardToProfileIdMap_managedCanNotForwardToAll()501     public void testGetCanForwardToProfileIdMap_managedCanNotForwardToAll() {
502         UserId currentUser = UserId.of(mManagedUser);
503         final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo2, mMockInfo3);
504 
505         if (SdkLevel.isAtLeastV()) {
506             initializeUserManagerState(currentUser,
507                     Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
508             when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
509                     PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
510                     mMockResolveInfoList);
511         } else {
512             initializeUserManagerState(currentUser,
513                     Lists.newArrayList(mSystemUser, mManagedUser));
514             when(mMockPackageManager.queryIntentActivities(mMockIntent,
515                     PackageManager.MATCH_DEFAULT_ONLY)).thenReturn(mMockResolveInfoList);
516         }
517 
518         Map<UserId, Boolean> expectedCanForwardToProfileIdMap = new HashMap<>();
519         expectedCanForwardToProfileIdMap.put(UserId.of(mSystemUser), false);
520         expectedCanForwardToProfileIdMap.put(UserId.of(mManagedUser), true);
521         if (SdkLevel.isAtLeastV()) {
522             expectedCanForwardToProfileIdMap.put(UserId.of(mPrivateUser), false);
523         }
524 
525         assertWithMessage("getCanForwardToProfileIdMap returns incorrect mappings")
526                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent))
527                 .isEqualTo(expectedCanForwardToProfileIdMap);
528     }
529 
530     @Test
testGetCanForwardToProfileIdMap_privateCanForwardToAll()531     public void testGetCanForwardToProfileIdMap_privateCanForwardToAll() {
532         if (!SdkLevel.isAtLeastV()) return;
533         UserId currentUser = UserId.of(mPrivateUser);
534         initializeUserManagerState(currentUser,
535                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
536         final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo1, mMockInfo2);
537         when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
538                 PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(mMockResolveInfoList);
539 
540         Map<UserId, Boolean> expectedCanForwardToProfileIdMap = new HashMap<>();
541         expectedCanForwardToProfileIdMap.put(UserId.of(mSystemUser), true);
542         expectedCanForwardToProfileIdMap.put(UserId.of(mManagedUser), true);
543         expectedCanForwardToProfileIdMap.put(UserId.of(mPrivateUser), true);
544 
545         assertWithMessage("getCanForwardToProfileIdMap returns incorrect mappings")
546                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent))
547                 .isEqualTo(expectedCanForwardToProfileIdMap);
548     }
549 
550     @Test
testGetCanForwardToProfileIdMap_privateCanNotForwardToManagedUser()551     public void testGetCanForwardToProfileIdMap_privateCanNotForwardToManagedUser() {
552         if (!SdkLevel.isAtLeastV()) return;
553         UserId currentUser = UserId.of(mPrivateUser);
554         initializeUserManagerState(currentUser,
555                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
556         final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo2, mMockInfo3);
557         when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
558                 PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(mMockResolveInfoList);
559 
560         Map<UserId, Boolean> expectedCanForwardToProfileIdMap = new HashMap<>();
561         expectedCanForwardToProfileIdMap.put(UserId.of(mSystemUser), true);
562         expectedCanForwardToProfileIdMap.put(UserId.of(mManagedUser), false);
563         expectedCanForwardToProfileIdMap.put(UserId.of(mPrivateUser), true);
564 
565         assertWithMessage("getCanForwardToProfileIdMap returns incorrect mappings")
566                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent))
567                 .isEqualTo(expectedCanForwardToProfileIdMap);
568     }
569 
570     @Test
testGetCanForwardToProfileIdMap_privateCanAlwaysForwardToSystemUser()571     public void testGetCanForwardToProfileIdMap_privateCanAlwaysForwardToSystemUser() {
572         if (!SdkLevel.isAtLeastV()) return;
573         UserId currentUser = UserId.of(mPrivateUser);
574         initializeUserManagerState(currentUser, Lists.newArrayList(mSystemUser, mPrivateUser));
575 
576         Map<UserId, Boolean> expectedCanForwardToProfileIdMap = new HashMap<>();
577         expectedCanForwardToProfileIdMap.put(UserId.of(mSystemUser), true);
578         expectedCanForwardToProfileIdMap.put(UserId.of(mPrivateUser), true);
579 
580         assertWithMessage("getCanForwardToProfileIdMap returns incorrect mappings")
581                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent))
582                 .isEqualTo(expectedCanForwardToProfileIdMap);
583     }
584 
585     @Test
testOnProfileStatusChange_anyIntentActionForManagedProfile()586     public void testOnProfileStatusChange_anyIntentActionForManagedProfile() {
587         if (!SdkLevel.isAtLeastV()) return;
588         UserId currentUser = UserId.of(mSystemUser);
589         initializeUserManagerState(currentUser,
590                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
591 
592         // UserManagerState#mUserId and UserManagerState#mCanForwardToProfileIdMap will empty
593         // by default if the getters of these member variables have not been called
594         List<UserId> userIdsBeforeIntent = new ArrayList<>(mUserManagerState.getUserIds());
595         Map<UserId, Boolean> canForwardToProfileIdMapBeforeIntent = new HashMap<>(
596                 mUserManagerState.getCanForwardToProfileIdMap(mMockIntent));
597 
598         String action = "any_intent";
599         mUserManagerState.onProfileActionStatusChange(action, UserId.of(mManagedUser));
600 
601         assertWithMessage("Unexpected changes to user id list on receiving intent: " + action)
602                 .that(mUserManagerState.getUserIds()).isEqualTo(userIdsBeforeIntent);
603         assertWithMessage(
604                 "Unexpected changes to canForwardToProfileIdMap on receiving intent: " + action)
605                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isEqualTo(
606                         canForwardToProfileIdMapBeforeIntent);
607     }
608 
609     @Test
testOnProfileStatusChange_actionProfileUnavailableForPrivateProfile()610     public void testOnProfileStatusChange_actionProfileUnavailableForPrivateProfile() {
611         if (!SdkLevel.isAtLeastV()) return;
612         UserId currentUser = UserId.of(mSystemUser);
613         UserId managedUser = UserId.of(mManagedUser);
614         UserId privateUser = UserId.of(mPrivateUser);
615         final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo1, mMockInfo2);
616         when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
617                 PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
618                 mMockResolveInfoList);
619         initializeUserManagerState(currentUser,
620                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
621 
622         // UserManagerState#mUserId and UserManagerState#mCanForwardToProfileIdMap will empty
623         // by default if the getters of these member variables have not been called
624         List<UserId> userIdsBeforeIntent = new ArrayList<>(mUserManagerState.getUserIds());
625         Map<UserId, Boolean> canForwardToProfileIdMapBeforeIntent = new HashMap<>(
626                 mUserManagerState.getCanForwardToProfileIdMap(mMockIntent));
627 
628         List<UserId> expectedUserIdsAfterIntent = Lists.newArrayList(currentUser, managedUser);
629 
630         String action = Intent.ACTION_PROFILE_UNAVAILABLE;
631         mUserManagerState.onProfileActionStatusChange(action, privateUser);
632 
633         assertWithMessage(
634                 "UserIds list should not be same before and after receiving intent: " + action)
635                 .that(mUserManagerState.getUserIds()).isNotEqualTo(userIdsBeforeIntent);
636         assertWithMessage("Unexpected changes to user id list on receiving intent: " + action)
637                 .that(mUserManagerState.getUserIds()).isEqualTo(expectedUserIdsAfterIntent);
638         assertWithMessage("CanForwardToLabelMap should be same before and after receiving intent: "
639                 + action)
640                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isEqualTo(
641                         canForwardToProfileIdMapBeforeIntent);
642     }
643 
644     @Test
testOnProfileStatusChange_actionProfileAvailable_profileInitialised()645     public void testOnProfileStatusChange_actionProfileAvailable_profileInitialised() {
646         if (!SdkLevel.isAtLeastV()) return;
647         UserId currentUser = UserId.of(mSystemUser);
648         UserId managedUser = UserId.of(mManagedUser);
649         UserId privateUser = UserId.of(mPrivateUser);
650         final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo1, mMockInfo2);
651         when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
652                 PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
653                 mMockResolveInfoList);
654         initializeUserManagerState(currentUser,
655                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
656 
657         // initialising the userIds list and canForwardToProfileIdMap
658         mUserManagerState.getUserIds();
659         mUserManagerState.getCanForwardToProfileIdMap(mMockIntent);
660 
661         // Making the private profile unavailable after it has been initialised
662         mUserManagerState.onProfileActionStatusChange(Intent.ACTION_PROFILE_UNAVAILABLE,
663                 privateUser);
664 
665         List<UserId> userIdsBeforeIntent = new ArrayList<>(mUserManagerState.getUserIds());
666         Map<UserId, Boolean> canForwardToProfileIdMapBeforeIntent = new HashMap<>(
667                 mUserManagerState.getCanForwardToProfileIdMap(mMockIntent));
668 
669         List<UserId> expectedUserIdsAfterIntent = Lists.newArrayList(currentUser, managedUser,
670                 privateUser);
671 
672         String action = Intent.ACTION_PROFILE_AVAILABLE;
673         mUserManagerState.onProfileActionStatusChange(action, privateUser);
674 
675         assertWithMessage(
676                 "UserIds list should not be same before and after receiving intent: " + action)
677                 .that(mUserManagerState.getUserIds()).isNotEqualTo(userIdsBeforeIntent);
678         assertWithMessage("Unexpected changes to user id list on receiving intent: " + action)
679                 .that(mUserManagerState.getUserIds()).isEqualTo(expectedUserIdsAfterIntent);
680         assertWithMessage("CanForwardToLabelMap should be same before and after receiving intent: "
681                 + action)
682                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isEqualTo(
683                         canForwardToProfileIdMapBeforeIntent);
684     }
685 
686     @Test
testOnProfileStatusChange_actionProfileAvailable_profileNotInitialised()687     public void testOnProfileStatusChange_actionProfileAvailable_profileNotInitialised() {
688         if (!SdkLevel.isAtLeastV()) return;
689         UserId currentUser = UserId.of(mSystemUser);
690         UserId managedUser = UserId.of(mManagedUser);
691         UserId privateUser = UserId.of(mPrivateUser);
692         final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo1, mMockInfo2);
693         when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
694                 PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
695                 mMockResolveInfoList);
696 
697         // Private user will not be initialised if it is in quiet mode
698         when(mMockUserManager.isQuietModeEnabled(mPrivateUser)).thenReturn(true);
699         initializeUserManagerState(currentUser,
700                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
701 
702         // UserManagerState#mUserId and UserManagerState#mCanForwardToProfileIdMap will be empty
703         // by default if the getters of these member variables have not been called
704         List<UserId> userIdsBeforeIntent = new ArrayList<>(mUserManagerState.getUserIds());
705         Map<UserId, Boolean> canForwardToProfileIdMapBeforeIntent = new HashMap<>(
706                 mUserManagerState.getCanForwardToProfileIdMap(mMockIntent));
707 
708         List<UserId> expectedUserIdsAfterIntent = Lists.newArrayList(currentUser, managedUser,
709                 privateUser);
710         Map<UserId, Boolean> expectedCanForwardToProfileIdMapAfterIntent = new HashMap<>();
711         expectedCanForwardToProfileIdMapAfterIntent.put(currentUser, true);
712         expectedCanForwardToProfileIdMapAfterIntent.put(managedUser, true);
713         expectedCanForwardToProfileIdMapAfterIntent.put(privateUser, true);
714 
715         String action = Intent.ACTION_PROFILE_AVAILABLE;
716         mUserManagerState.onProfileActionStatusChange(action, privateUser);
717 
718         assertWithMessage(
719                 "UserIds list should not be same before and after receiving intent: " + action)
720                 .that(mUserManagerState.getUserIds()).isNotEqualTo(userIdsBeforeIntent);
721         assertWithMessage("Unexpected changes to user id list on receiving intent: " + action)
722                 .that(mUserManagerState.getUserIds()).isEqualTo(expectedUserIdsAfterIntent);
723         assertWithMessage(
724                 "CanForwardToLabelMap should not be same before and after receiving intent: "
725                         + action)
726                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isNotEqualTo(
727                         canForwardToProfileIdMapBeforeIntent);
728         assertWithMessage(
729                 "Unexpected changes to canForwardToProfileIdMap on receiving intent: " + action)
730                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isEqualTo(
731                         expectedCanForwardToProfileIdMapAfterIntent);
732     }
733 
734     @Test
testGetUserIdToLabelMap_systemUserAndManagedUser_PreV()735     public void testGetUserIdToLabelMap_systemUserAndManagedUser_PreV() {
736         if (SdkLevel.isAtLeastV()) return;
737         UserId currentUser = UserId.of(mSystemUser);
738         initializeUserManagerState(currentUser,
739                 Lists.newArrayList(mSystemUser, mManagedUser));
740         if (SdkLevel.isAtLeastT()) {
741             DevicePolicyResourcesManager devicePolicyResourcesManager = mock(
742                     DevicePolicyResourcesManager.class);
743             when(mDevicePolicyManager.getResources()).thenReturn(devicePolicyResourcesManager);
744             when(devicePolicyResourcesManager.getString(eq(PERSONAL_TAB), any())).thenReturn(
745                     PERSONAL);
746             when(devicePolicyResourcesManager.getString(eq(WORK_TAB), any())).thenReturn(WORK);
747         }
748 
749         Map<UserId, String> userIdToLabelMap = mUserManagerState.getUserIdToLabelMap();
750 
751         assertWithMessage("Incorrect label returned for user id " + mSystemUser)
752                 .that(userIdToLabelMap.get(UserId.of(mSystemUser))).isEqualTo(PERSONAL);
753         assertWithMessage("Incorrect label returned for user id " + mManagedUser)
754                 .that(userIdToLabelMap.get(UserId.of(mManagedUser))).isEqualTo(WORK);
755     }
756 
757     @Test
testGetUserIdToLabelMap_systemUserManagedUserPrivateUser_PostV()758     public void testGetUserIdToLabelMap_systemUserManagedUserPrivateUser_PostV() {
759         if (!SdkLevel.isAtLeastV()) return;
760         UserId currentUser = UserId.of(mSystemUser);
761         initializeUserManagerState(currentUser,
762                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
763         if (SdkLevel.isAtLeastT()) {
764             DevicePolicyResourcesManager devicePolicyResourcesManager = mock(
765                     DevicePolicyResourcesManager.class);
766             when(mDevicePolicyManager.getResources()).thenReturn(devicePolicyResourcesManager);
767             when(devicePolicyResourcesManager.getString(eq(PERSONAL_TAB), any())).thenReturn(
768                     PERSONAL);
769         }
770         UserManager managedUserManager = getUserManagerForManagedUser();
771         UserManager privateUserManager = getUserManagerForPrivateUser();
772         when(managedUserManager.getProfileLabel()).thenReturn(WORK);
773         when(privateUserManager.getProfileLabel()).thenReturn(PRIVATE);
774 
775         Map<UserId, String> userIdToLabelMap = mUserManagerState.getUserIdToLabelMap();
776 
777         assertWithMessage("Incorrect label returned for user id " + mSystemUser)
778                 .that(userIdToLabelMap.get(UserId.of(mSystemUser))).isEqualTo(PERSONAL);
779         assertWithMessage("Incorrect label returned for user id " + mManagedUser)
780                 .that(userIdToLabelMap.get(UserId.of(mManagedUser))).isEqualTo(WORK);
781         assertWithMessage("Incorrect label returned for user id " + mPrivateUser)
782                 .that(userIdToLabelMap.get(UserId.of(mPrivateUser))).isEqualTo(PRIVATE);
783     }
784 
785     @Test
testGetUserIdToBadgeMap_systemUserManagedUser_PreV()786     public void testGetUserIdToBadgeMap_systemUserManagedUser_PreV() {
787         if (SdkLevel.isAtLeastV()) return;
788         UserId currentUser = UserId.of(mSystemUser);
789         initializeUserManagerState(currentUser,
790                 Lists.newArrayList(mSystemUser, mManagedUser));
791         Drawable workBadge = mock(Drawable.class);
792         Resources resources = mock(Resources.class);
793         when(mMockContext.getResources()).thenReturn(resources);
794         when(mMockContext.getDrawable(R.drawable.ic_briefcase)).thenReturn(workBadge);
795         if (SdkLevel.isAtLeastT()) {
796             DevicePolicyResourcesManager devicePolicyResourcesManager = mock(
797                     DevicePolicyResourcesManager.class);
798             when(mDevicePolicyManager.getResources()).thenReturn(devicePolicyResourcesManager);
799             when(devicePolicyResourcesManager.getDrawable(eq(WORK_PROFILE_ICON), eq(SOLID_COLORED),
800                     any())).thenReturn(workBadge);
801         }
802 
803         Map<UserId, Drawable> userIdToBadgeMap = mUserManagerState.getUserIdToBadgeMap();
804 
805         assertWithMessage("There should be no badge present for personal user")
806                 .that(userIdToBadgeMap.containsKey(UserId.of(mSystemUser))).isFalse();
807         assertWithMessage("Incorrect badge returned for user id " + mManagedUser)
808                 .that(userIdToBadgeMap.get(UserId.of(mManagedUser))).isEqualTo(workBadge);
809     }
810 
811     @Test
testGetUserIdToBadgeMap_systemUserManagedUserPrivateUser_PostV()812     public void testGetUserIdToBadgeMap_systemUserManagedUserPrivateUser_PostV() {
813         if (!SdkLevel.isAtLeastV()) return;
814         UserId currentUser = UserId.of(mSystemUser);
815         initializeUserManagerState(currentUser,
816                 Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
817         Drawable workBadge = mock(Drawable.class);
818         Drawable privateBadge = mock(Drawable.class);
819         UserManager managedUserManager = getUserManagerForManagedUser();
820         UserManager privateUserManager = getUserManagerForPrivateUser();
821         when(managedUserManager.getUserBadge()).thenReturn(workBadge);
822         when(privateUserManager.getUserBadge()).thenReturn(privateBadge);
823 
824         Map<UserId, Drawable> userIdToBadgeMap = mUserManagerState.getUserIdToBadgeMap();
825 
826         assertWithMessage("There should be no badge present for personal user")
827                 .that(userIdToBadgeMap.get(UserId.of(mSystemUser))).isNull();
828         assertWithMessage("Incorrect badge returned for user id " + mManagedUser)
829                 .that(userIdToBadgeMap.get(UserId.of(mManagedUser))).isEqualTo(workBadge);
830         assertWithMessage("Incorrect badge returned for user id " + mPrivateUser)
831                 .that(userIdToBadgeMap.get(UserId.of(mPrivateUser))).isEqualTo(privateBadge);
832     }
833 
initializeUserManagerState(UserId current, List<UserHandle> usersOnDevice)834     private void initializeUserManagerState(UserId current, List<UserHandle> usersOnDevice) {
835         when(mMockUserManager.getUserProfiles()).thenReturn(usersOnDevice);
836         TestConfigStore testConfigStore = new TestConfigStore();
837         testConfigStore.enablePrivateSpaceInPhotoPicker();
838         mUserManagerState = new UserManagerState.RuntimeUserManagerState(mMockContext, current,
839                 true, testConfigStore);
840     }
841 
getUserManagerForManagedUser()842     private UserManager getUserManagerForManagedUser() {
843         Context managedUserContext = mock(Context.class);
844         when(mMockContext.createContextAsUser(mManagedUser, 0)).thenReturn(managedUserContext);
845         UserManager managedUserManager = mock(UserManager.class);
846         when(managedUserContext.getSystemServiceName(UserManager.class))
847                 .thenReturn("managedUserManager");
848         when(managedUserContext.getSystemService(UserManager.class)).thenReturn(managedUserManager);
849         return managedUserManager;
850     }
851 
getUserManagerForPrivateUser()852     private UserManager getUserManagerForPrivateUser() {
853         Context privateUserContext = mock(Context.class);
854         when(mMockContext.createContextAsUser(mPrivateUser, 0)).thenReturn(privateUserContext);
855         UserManager privateUserManager = mock(UserManager.class);
856         when(privateUserContext.getSystemServiceName(UserManager.class))
857                 .thenReturn("privateUserManager");
858         when(privateUserContext.getSystemService(UserManager.class)).thenReturn(privateUserManager);
859         return privateUserManager;
860     }
861 }
862