Fix wrong user in personal profile

If the user is profile, it should use the current
user's profile parent to create the context instead
of always using the primary user(system user).

Bug: 264867495
Test: Manual test with work profile enabled, multi user flow, etc
Test: make RunSettingsRoboTests ROBOTEST_FILTER=AvailableVirtualKeyboardFragmentTest
Change-Id: I66d1462441359c40dda2a38c7534500791db9c83

src/com/android/settings/inputmethod/AvailableVirtualKeyboardFragment.java

@@ -55,7 +55,10 @@ public class AvailableVirtualKeyboardFragment extends DashboardFragment
 
     @VisibleForTesting
     InputMethodSettingValuesWrapper mInputMethodSettingValues;
-    private Context mUserAwareContext;
+
+    @VisibleForTesting
+    Context mUserAwareContext;
+
     private int mUserId;
 
     @Override
@@ -82,9 +85,16 @@ public class AvailableVirtualKeyboardFragment extends DashboardFragment
                 break;
             }
             case ProfileSelectFragment.ProfileType.PERSONAL: {
-                final UserHandle primaryUser = userManager.getPrimaryUser().getUserHandle();
-                newUserId = primaryUser.getIdentifier();
-                newUserAwareContext = context.createContextAsUser(primaryUser, 0);
+                // Use the parent user of the current user if the current user is profile.
+                final UserHandle currentUser = UserHandle.of(currentUserId);
+                final UserHandle userProfileParent = userManager.getProfileParent(currentUser);
+                if (userProfileParent != null) {
+                    newUserId = userProfileParent.getIdentifier();
+                    newUserAwareContext = context.createContextAsUser(userProfileParent, 0);
+                } else {
+                    newUserId = currentUserId;
+                    newUserAwareContext = context;
+                }
                 break;
             }
             default:

tests/robotests/src/com/android/settings/inputmethod/AvailableVirtualKeyboardFragmentTest.java

@@ -27,6 +27,7 @@ import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.mock;
 
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
@@ -34,6 +35,7 @@ import android.content.pm.ResolveInfo;
 import android.content.pm.ServiceInfo;
 import android.os.Bundle;
 import android.os.UserHandle;
+import android.os.UserManager;
 import android.provider.SearchIndexableResource;
 import android.view.inputmethod.InputMethodInfo;
 import android.view.inputmethod.InputMethodManager;
@@ -46,7 +48,6 @@ import com.android.settings.R;
 import com.android.settings.dashboard.profileselector.ProfileSelectFragment;
 import com.android.settings.testutils.shadow.ShadowInputMethodManagerWithMethodList;
 import com.android.settings.testutils.shadow.ShadowSecureSettings;
-import com.android.settings.testutils.shadow.ShadowUserManager;
 import com.android.settingslib.inputmethod.InputMethodPreference;
 import com.android.settingslib.inputmethod.InputMethodSettingValuesWrapper;
 
@@ -66,14 +67,15 @@ import java.util.List;
 @RunWith(RobolectricTestRunner.class)
 @Config(shadows = {
         ShadowSecureSettings.class,
-        ShadowInputMethodManagerWithMethodList.class,
-        ShadowUserManager.class
+        ShadowInputMethodManagerWithMethodList.class
 })
 public class AvailableVirtualKeyboardFragmentTest {
 
     @Mock
     private InputMethodManager mInputMethodManager;
     @Mock
+    private UserManager mUserManager;
+    @Mock
     private InputMethodSettingValuesWrapper mValuesWrapper;
     @Mock
     private PreferenceScreen mPreferenceScreen;
@@ -92,6 +94,27 @@ public class AvailableVirtualKeyboardFragmentTest {
         initMock();
     }
 
+    @Test
+    public void onAttachPersonalProfile_noProfileParent() {
+        doReturn(null).when(mUserManager).getProfileParent(any(UserHandle.class));
+
+        mFragment.onAttach(mContext);
+
+        assertThat(mFragment.mUserAwareContext).isEqualTo(mContext);
+    }
+
+    @Test
+    public void onAttachPersonalProfile_hasProfileParent() {
+        final UserHandle profileParent = new UserHandle(0);
+        final Context mockContext = mock(Context.class);
+        doReturn(profileParent).when(mUserManager).getProfileParent(any(UserHandle.class));
+        doReturn(mockContext).when(mContext).createContextAsUser(any(UserHandle.class), anyInt());
+
+        mFragment.onAttach(mContext);
+
+        assertThat(mFragment.mUserAwareContext).isEqualTo(mockContext);
+    }
+
     @Test
     public void onCreatePreferences_shouldAddResource() {
         mFragment.onAttach(mContext);
@@ -175,7 +198,7 @@ public class AvailableVirtualKeyboardFragmentTest {
         when(mFragment.getPreferenceScreen()).thenReturn(mPreferenceScreen);
         when(mPreferenceManager.getContext()).thenReturn(mContext);
         when(mContext.getSystemService(InputMethodManager.class)).thenReturn(mInputMethodManager);
-        doReturn(mContext).when(mContext).createContextAsUser(any(UserHandle.class), anyInt());
+        when(mContext.getSystemService(UserManager.class)).thenReturn(mUserManager);
     }
 
     private List<InputMethodInfo> createFakeInputMethodInfoList(final String name, int num) {