Fix postsubmit test failures

Test: atest ProfileSelectFragmentTest
Bug: 314371622
Change-Id: Id31518fb92491e5a434cc4a3a95b6353d58a382f

src/com/android/settings/dashboard/profileselector/ProfileSelectFragment.java

@@ -41,6 +41,8 @@ import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import androidx.fragment.app.Fragment;
+import androidx.fragment.app.FragmentManager;
+import androidx.lifecycle.Lifecycle;
 import androidx.recyclerview.widget.RecyclerView;
 import androidx.viewpager2.adapter.FragmentStateAdapter;
 import androidx.viewpager2.widget.ViewPager2;
@@ -357,6 +359,11 @@ public abstract class ProfileSelectFragment extends DashboardFragment {
         return fragment;
     }
 
+    @VisibleForTesting
+    void setViewPager(ViewPager2 viewPager) {
+        mViewPager = viewPager;
+    }
+
     interface FragmentConstructor {
         Fragment constructAndGetFragment();
     }
@@ -376,6 +383,15 @@ public abstract class ProfileSelectFragment extends DashboardFragment {
             mChildFragments = fragment.getFragments();
         }
 
+        @VisibleForTesting
+        ViewPagerAdapter(
+                @NonNull FragmentManager fragmentManager,
+                @NonNull Lifecycle lifecycle,
+                ProfileSelectFragment profileSelectFragment) {
+            super(fragmentManager, lifecycle);
+            mChildFragments = profileSelectFragment.getFragments();
+        }
+
         @Override
         public Fragment createFragment(int position) {
             return mChildFragments[position];
@@ -386,7 +402,8 @@ public abstract class ProfileSelectFragment extends DashboardFragment {
             return mChildFragments.length;
         }
 
-        private int getTabForPosition(int position) {
+        @VisibleForTesting
+        int getTabForPosition(int position) {
             if (position >= mChildFragments.length) {
                 Log.e(TAG, "tab requested for out of bound position " + position);
                 return PERSONAL_TAB;

tests/robotests/src/com/android/settings/dashboard/profileselector/ProfileSelectFragmentTest.java

@@ -39,17 +39,24 @@ import android.os.Flags;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.util.ArraySet;
 
+import androidx.annotation.NonNull;
 import androidx.fragment.app.Fragment;
 import androidx.fragment.app.FragmentActivity;
+import androidx.fragment.app.FragmentManager;
+import androidx.fragment.app.FragmentTransaction;
+import androidx.lifecycle.Lifecycle;
+import androidx.viewpager2.widget.ViewPager2;
 
 import com.android.settings.SettingsActivity;
 import com.android.settings.SettingsPreferenceFragmentTest;
+import com.android.settings.dashboard.profileselector.ProfileSelectFragment.ViewPagerAdapter;
 import com.android.settings.testutils.shadow.ShadowUserManager;
 
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
@@ -75,6 +82,9 @@ public class ProfileSelectFragmentTest {
     private TestProfileSelectFragment mFragment;
     private FragmentActivity mActivity;
     private ShadowUserManager mUserManager;
+    @Mock private FragmentManager mFragmentManager;
+    @Mock private Lifecycle mLifecycle;
+    @Mock private FragmentTransaction mFragmentTransaction;
     @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
     @Before
@@ -97,6 +107,16 @@ public class ProfileSelectFragmentTest {
     public void getTabId_setArgumentWork_setCorrectTab() {
         final Bundle bundle = new Bundle();
         bundle.putInt(SettingsActivity.EXTRA_SHOW_FRAGMENT_TAB, WORK_TAB);
+        ViewPager2 viewPager = new ViewPager2(mContext);
+        TestProfileSelectFragment profileSelectFragment = new TestProfileSelectFragment();
+        ViewPagerAdapter viewPagerAdapter =
+                new TestViewPagerAdapter(mFragmentManager, mLifecycle, profileSelectFragment);
+
+        when(mFragmentManager.beginTransaction()).thenReturn(mFragmentTransaction);
+        viewPager.setAdapter(viewPagerAdapter);
+        mFragment.setViewPager(viewPager);
+        profileSelectFragment.setViewPager(viewPager);
+        mFragmentManager.beginTransaction().add(profileSelectFragment, "tag");
 
         assertThat(mFragment.getTabId(mActivity, bundle)).isEqualTo(WORK_TAB);
     }
@@ -105,6 +125,16 @@ public class ProfileSelectFragmentTest {
     public void getTabId_setArgumentPrivate_setCorrectTab() {
         final Bundle bundle = new Bundle();
         bundle.putInt(SettingsActivity.EXTRA_SHOW_FRAGMENT_TAB, PRIVATE_TAB);
+        ViewPager2 viewPager = new ViewPager2(mContext);
+        TestProfileSelectFragment profileSelectFragment = new TestProfileSelectFragment();
+        ViewPagerAdapter viewPagerAdapter =
+                new TestViewPagerAdapter(mFragmentManager, mLifecycle, profileSelectFragment);
+
+        when(mFragmentManager.beginTransaction()).thenReturn(mFragmentTransaction);
+        viewPager.setAdapter(viewPagerAdapter);
+        mFragment.setViewPager(viewPager);
+        profileSelectFragment.setViewPager(viewPager);
+        mFragmentManager.beginTransaction().add(profileSelectFragment, "tag");
 
         assertThat(mFragment.getTabId(mActivity, bundle)).isEqualTo(PRIVATE_TAB);
     }
@@ -113,6 +143,16 @@ public class ProfileSelectFragmentTest {
     public void getTabId_setArgumentPersonal_setCorrectTab() {
         final Bundle bundle = new Bundle();
         bundle.putInt(SettingsActivity.EXTRA_SHOW_FRAGMENT_TAB, PERSONAL_TAB);
+        ViewPager2 viewPager = new ViewPager2(mContext);
+        TestProfileSelectFragment profileSelectFragment = new TestProfileSelectFragment();
+        ViewPagerAdapter viewPagerAdapter =
+                new TestViewPagerAdapter(mFragmentManager, mLifecycle, profileSelectFragment);
+
+        when(mFragmentManager.beginTransaction()).thenReturn(mFragmentTransaction);
+        viewPager.setAdapter(viewPagerAdapter);
+        mFragment.setViewPager(viewPager);
+        profileSelectFragment.setViewPager(viewPager);
+        mFragmentManager.beginTransaction().add(profileSelectFragment, "tag");
 
         assertThat(mFragment.getTabId(mActivity, bundle)).isEqualTo(PERSONAL_TAB);
     }
@@ -310,4 +350,18 @@ public class ProfileSelectFragmentTest {
             };
         }
     }
+
+    static class TestViewPagerAdapter extends ViewPagerAdapter {
+        TestViewPagerAdapter(
+                @NonNull FragmentManager fragmentManager,
+                @NonNull Lifecycle lifecycle,
+                ProfileSelectFragment profileSelectFragment) {
+            super(fragmentManager, lifecycle, profileSelectFragment);
+        }
+
+        @Override
+        int getTabForPosition(int position) {
+            return position;
+        }
+    }
 }