From 8ca4f47292847ac48d7aacd09c7c64e99868324f Mon Sep 17 00:00:00 2001
From: Manish Singh <psych@google.com>
Date: Tue, 5 Mar 2024 16:28:33 +0000
Subject: [PATCH] Return the correct position for the profile tab

The existing code returns the fixed tab values, which doesn't work
now that the ordering is determined by the user manager's ordering.

Bug: 302278487
Test: manual
Test: atest ProfileSelectFragmentTest
Change-Id: I87b393b8e12e21dc1b38ee75bb43ca9133785c81
---
 .../ProfileSelectFragment.java                | 18 +++++++++++-
 .../ProfileSelectFragmentTest.java            | 29 +++++++++++++++----
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/src/com/android/settings/dashboard/profileselector/ProfileSelectFragment.java b/src/com/android/settings/dashboard/profileselector/ProfileSelectFragment.java
index 8279588cd1d..53a55542848 100644
--- a/src/com/android/settings/dashboard/profileselector/ProfileSelectFragment.java
+++ b/src/com/android/settings/dashboard/profileselector/ProfileSelectFragment.java
@@ -230,7 +230,8 @@ public abstract class ProfileSelectFragment extends DashboardFragment {
         if (bundle != null) {
             final int extraTab = bundle.getInt(SettingsActivity.EXTRA_SHOW_FRAGMENT_TAB, -1);
             if (extraTab != -1) {
-                return ((ViewPagerAdapter) mViewPager.getAdapter()).getTabForPosition(extraTab);
+                return ((ViewPagerAdapter) mViewPager.getAdapter())
+                        .getPositionForProfileTab(extraTab);
             }
             final int userId = bundle.getInt(EXTRA_USER_ID, UserHandle.SYSTEM.getIdentifier());
             final boolean isWorkProfile = UserManager.get(activity).isManagedProfile(userId);
@@ -410,7 +411,22 @@ public abstract class ProfileSelectFragment extends DashboardFragment {
             }
             @ProfileType
             int profileType = mChildFragments[position].getArguments().getInt(EXTRA_PROFILE);
+            return profileTypeToTab(profileType);
+        }
 
+        private int getPositionForProfileTab(int profileTab) {
+            for (int i = 0; i < mChildFragments.length; ++i) {
+                Bundle arguments = mChildFragments[i].getArguments();
+                if (arguments != null
+                        && profileTypeToTab(arguments.getInt(EXTRA_PROFILE)) == profileTab) {
+                    return i;
+                }
+            }
+            Log.e(TAG, "position requested for an unknown profile tab " + profileTab);
+            return 0;
+        }
+
+        private int profileTypeToTab(@ProfileType int profileType) {
             if (profileType == ProfileType.WORK) {
                 return WORK_TAB;
             }
diff --git a/tests/robotests/src/com/android/settings/dashboard/profileselector/ProfileSelectFragmentTest.java b/tests/robotests/src/com/android/settings/dashboard/profileselector/ProfileSelectFragmentTest.java
index 17e0d1ca63c..b6af0f908cc 100644
--- a/tests/robotests/src/com/android/settings/dashboard/profileselector/ProfileSelectFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/dashboard/profileselector/ProfileSelectFragmentTest.java
@@ -118,7 +118,9 @@ public class ProfileSelectFragmentTest {
         profileSelectFragment.setViewPager(viewPager);
         mFragmentManager.beginTransaction().add(profileSelectFragment, "tag");
 
-        assertThat(mFragment.getTabId(mActivity, bundle)).isEqualTo(WORK_TAB);
+        // The expected position '2' comes from the order in which fragments are added in
+        // TestProfileSelectFragment#getFragments()
+        assertThat(mFragment.getTabId(mActivity, bundle)).isEqualTo(2);
     }
 
     @Test
@@ -136,7 +138,9 @@ public class ProfileSelectFragmentTest {
         profileSelectFragment.setViewPager(viewPager);
         mFragmentManager.beginTransaction().add(profileSelectFragment, "tag");
 
-        assertThat(mFragment.getTabId(mActivity, bundle)).isEqualTo(PRIVATE_TAB);
+        // The expected position '1' comes from the order in which fragments are added in
+        // TestProfileSelectFragment#getFragments()
+        assertThat(mFragment.getTabId(mActivity, bundle)).isEqualTo(1);
     }
 
     @Test
@@ -343,10 +347,25 @@ public class ProfileSelectFragmentTest {
 
         @Override
         public Fragment[] getFragments() {
+            Fragment personalFragment = new SettingsPreferenceFragmentTest.TestFragment();
+            Bundle personalBundle = new Bundle();
+            personalBundle.putInt(EXTRA_PROFILE, ProfileType.PERSONAL);
+            personalFragment.setArguments(personalBundle);
+
+            Fragment workFragment = new SettingsPreferenceFragmentTest.TestFragment();
+            Bundle workBundle = new Bundle();
+            workBundle.putInt(EXTRA_PROFILE, ProfileType.WORK);
+            workFragment.setArguments(workBundle);
+
+            Fragment privateFragment = new SettingsPreferenceFragmentTest.TestFragment();
+            Bundle privateBundle = new Bundle();
+            privateBundle.putInt(EXTRA_PROFILE, ProfileType.PRIVATE);
+            privateFragment.setArguments(privateBundle);
+
             return new Fragment[]{
-                    new SettingsPreferenceFragmentTest.TestFragment(), //0
-                    new SettingsPreferenceFragmentTest.TestFragment(),
-                    new SettingsPreferenceFragmentTest.TestFragment()
+                    personalFragment, //0
+                    privateFragment,
+                    workFragment
             };
         }
     }