Aggregate TabModel and List<Tab> before detaching anything

Relying on tabProvider.get() as an iterator causes an infinite loop on
tablets. As an alternative, I aggregate all Tabs under each TabModel
and iterate over them explicitly.

Bug: 1046308
Change-Id: I3f096a0549ef9a7d463acc4eebbed46f7c9aac99
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2033703
Commit-Queue: Brandon Wylie <wylieb@chromium.org>
Reviewed-by: Yusuf Ozuysal <yusufo@chromium.org>
Cr-Commit-Position: refs/heads/master@{#739124}

chrome/android/java/src/org/chromium/chrome/browser/night_mode/NightModeReparentingController.java

@@ -14,10 +14,16 @@ import org.chromium.chrome.browser.tab.TabLaunchType;
 import org.chromium.chrome.browser.tab_activity_glue.ReparentingTask;
 import org.chromium.chrome.browser.tabmodel.AsyncTabParams;
 import org.chromium.chrome.browser.tabmodel.AsyncTabParamsManager;
+import org.chromium.chrome.browser.tabmodel.TabList;
 import org.chromium.chrome.browser.tabmodel.TabModel;
 import org.chromium.chrome.browser.tabmodel.TabModelSelector;
 import org.chromium.chrome.browser.tabmodel.TabReparentingParams;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
 // TODO(wylieb): Write unittests for this class.
 /** Controls the reparenting of tabs when the theme is swapped. */
 public class NightModeReparentingController implements NightModeStateProvider.Observer {
@@ -93,31 +99,55 @@ public class NightModeReparentingController implements NightModeStateProvider.Ob
 
     @Override
     public void onNightModeStateChanged() {
-        ActivityTabProvider tabProvider = mDelegate.getActivityTabProvider();
-
-        boolean isForegroundTab = true;
-        while (tabProvider.get() != null) {
-            Tab tabToDetach = tabProvider.get();
-            TabModel tabModel = mDelegate.getTabModelSelector().getModel(tabToDetach.isIncognito());
-            if (tabModel == null) continue;
-
-            TabReparentingParams params = new TabReparentingParams(tabToDetach, null, null);
-            params.setFromNightModeReparenting(true);
-
-            // Only the first tab is considered foreground and that is the only one for which
-            // we need an index. It will be reattached at the end. All remaining tabs will be
-            // reattached in reverse order, therefore we can ignore the index.
-            if (isForegroundTab) {
-                params.setIsForegroundTab(true);
-                params.setTabIndex(tabModel.indexOf(tabToDetach));
-                isForegroundTab = false;
-            } else {
-                params.setIsForegroundTab(false);
+        Tab foregroundTab = mDelegate.getActivityTabProvider().get();
+        Map<TabModel, List<Tab>> tabsAndModels = getTabModelAndTabs();
+
+        for (final TabModel model : tabsAndModels.keySet()) {
+            List<Tab> tabs = tabsAndModels.get(model);
+            for (int i = 0; i < tabs.size(); i++) {
+                Tab tab = tabs.get(i);
+                TabReparentingParams params = new TabReparentingParams(tab, null, null);
+                params.setFromNightModeReparenting(true);
+
+                // Only the first tab is considered foreground and that is the only one for which
+                // we need an index. It will be reattached at the end. All remaining tabs will be
+                // reattached in reverse order, therefore we can ignore the index.
+                if (tab == foregroundTab) {
+                    params.setIsForegroundTab(true);
+                    params.setTabIndex(i);
+                } else {
+                    params.setIsForegroundTab(false);
+                }
+
+                AsyncTabParamsManager.add(tab.getId(), params);
+                model.removeTab(tab);
+                ReparentingTask.from(tab).detach();
             }
+        }
+    }
+
+    /** @return Mapping from TabModel to all the Tabs in that model. */
+    Map<TabModel, List<Tab>> getTabModelAndTabs() {
+        Map<TabModel, List<Tab>> tabsAndModels = new HashMap<>();
 
-            AsyncTabParamsManager.add(tabToDetach.getId(), params);
-            tabModel.removeTab(tabToDetach);
-            ReparentingTask.from(tabToDetach).detach();
+        TabModel model = mDelegate.getTabModelSelector().getModel(false);
+        tabsAndModels.put(model, getTabsForModel(model));
+        model = mDelegate.getTabModelSelector().getModel(true);
+        tabsAndModels.put(model, getTabsForModel(model));
+
+        return tabsAndModels;
+    }
+
+    /** @return All of a model's Tabs as a List. */
+    List<Tab> getTabsForModel(TabModel model) {
+        List<Tab> tabs = new ArrayList<>();
+        if (model == null) return tabs;
+
+        TabList tabList = model.getComprehensiveModel();
+        for (int i = 0; i < tabList.getCount(); i++) {
+            tabs.add(tabList.getTabAt(i));
         }
+
+        return tabs;
     }
 }

chrome/android/junit/src/org/chromium/chrome/browser/night_mode/NightModeReparentingControllerTest.java

@@ -16,8 +16,6 @@ import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
-import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
@@ -27,16 +25,15 @@ import org.chromium.base.UserDataHost;
 import org.chromium.base.test.BaseRobolectricTestRunner;
 import org.chromium.chrome.browser.ActivityTabProvider;
 import org.chromium.chrome.browser.tab.Tab;
+import org.chromium.chrome.browser.tab.TabLaunchType;
 import org.chromium.chrome.browser.tab_activity_glue.ReparentingTask;
 import org.chromium.chrome.browser.tabmodel.AsyncTabParams;
 import org.chromium.chrome.browser.tabmodel.AsyncTabParamsManager;
-import org.chromium.chrome.browser.tabmodel.TabModel;
 import org.chromium.chrome.browser.tabmodel.TabModelSelector;
 import org.chromium.chrome.browser.tabmodel.TabReparentingParams;
+import org.chromium.chrome.test.util.browser.tabmodel.MockTabModel;
 
-import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
 /**
@@ -54,7 +51,7 @@ public class NightModeReparentingControllerTest {
             if (mActivityTabProvider == null) {
                 // setup
                 mActivityTabProvider = Mockito.mock(ActivityTabProvider.class);
-                doAnswer(invocation -> getNextTabFromList()).when(mActivityTabProvider).get();
+                doAnswer(invocation -> getForegroundTab()).when(mActivityTabProvider).get();
             }
             return mActivityTabProvider;
         }
@@ -64,29 +61,24 @@ public class NightModeReparentingControllerTest {
             if (mTabModelSelector == null) {
                 // setup
                 mTabModelSelector = Mockito.mock(TabModelSelector.class);
-                doAnswer(invocation -> getCurrentTabModel())
-                        .when(mTabModelSelector)
-                        .getCurrentModel();
-                doAnswer(invocation -> getCurrentTabModel())
-                        .when(mTabModelSelector)
-                        .getModel(anyBoolean());
+
+                doReturn(mTabModel).when(mTabModelSelector).getModel(false);
+                doReturn(mIncognitoTabModel).when(mTabModelSelector).getModel(true);
             }
 
             return mTabModelSelector;
         }
     }
 
-    @Mock
-    TabModel mTabModel;
     @Mock
     ReparentingTask mTask;
     @Mock
     ReparentingTask.Delegate mReparentingTaskDelegate;
-    @Captor
-    ArgumentCaptor<Tab> mTabCaptor;
 
-    List<Tab> mTabs = new ArrayList<>();
+    MockTabModel mTabModel;
+    MockTabModel mIncognitoTabModel;
     Map<Tab, Integer> mTabIndexMapping = new HashMap<>();
+    Tab mForegroundTab;
 
     NightModeReparentingController mController;
     NightModeReparentingController.Delegate mDelegate;
@@ -95,16 +87,8 @@ public class NightModeReparentingControllerTest {
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
-        doAnswer(invocation -> {
-            removeTab(mTabCaptor.getValue());
-            return null;
-        })
-                .when(mTabModel)
-                .removeTab(mTabCaptor.capture());
-
-        doAnswer(invocation -> getTabIndex(mTabCaptor.getValue()))
-                .when(mTabModel)
-                .indexOf(mTabCaptor.capture());
+        mTabModel = new MockTabModel(false, null);
+        mIncognitoTabModel = new MockTabModel(true, null);
 
         mDelegate = new FakeNightModeReparentingDelegate();
         mController = new NightModeReparentingController(mDelegate, mReparentingTaskDelegate);
@@ -112,14 +96,14 @@ public class NightModeReparentingControllerTest {
 
     @After
     public void tearDown() {
+        mForegroundTab = null;
         AsyncTabParamsManager.getAsyncTabParams().clear();
-        mTabs.clear();
         mTabIndexMapping.clear();
     }
 
     @Test
     public void testReparenting_singleTab() {
-        createAndAddMockTab(1, 0);
+        mForegroundTab = createAndAddMockTab(1, false);
         mController.onNightModeStateChanged();
 
         AsyncTabParams params = AsyncTabParamsManager.getAsyncTabParams().get(1);
@@ -143,7 +127,7 @@ public class NightModeReparentingControllerTest {
 
     @Test
     public void testReparenting_singleTab_currentModelNullOnStart() {
-        createAndAddMockTab(1, 0);
+        mForegroundTab = createAndAddMockTab(1, false);
         mController.onNightModeStateChanged();
 
         doReturn(null).when(mDelegate.getTabModelSelector()).getModel(anyBoolean());
@@ -155,8 +139,8 @@ public class NightModeReparentingControllerTest {
 
     @Test
     public void testReparenting_multipleTabs_onlyOneIsReparented() {
-        createAndAddMockTab(1, 0);
-        createAndAddMockTab(2, 1);
+        mForegroundTab = createAndAddMockTab(1, false);
+        createAndAddMockTab(2, false);
         mController.onNightModeStateChanged();
 
         TabReparentingParams trp =
@@ -189,11 +173,11 @@ public class NightModeReparentingControllerTest {
 
     @Test
     public void testReparenting_twoTabsOutOfOrder() {
-        createAndAddMockTab(1, 1);
-        createAndAddMockTab(2, 0);
+        createAndAddMockTab(1, false);
+        mForegroundTab = createAndAddMockTab(2, false);
         mController.onNightModeStateChanged();
 
-        AsyncTabParams params = AsyncTabParamsManager.getAsyncTabParams().get(1);
+        AsyncTabParams params = AsyncTabParamsManager.getAsyncTabParams().get(2);
         Assert.assertNotNull(params);
         Assert.assertTrue(params instanceof TabReparentingParams);
 
@@ -206,7 +190,34 @@ public class NightModeReparentingControllerTest {
 
         Tab tab = trp.getTabToReparent();
         Assert.assertNotNull(tab);
-        Assert.assertEquals(1, tab.getId());
+        Assert.assertEquals(2, tab.getId());
+
+        verify(mTask, times(2)).detach();
+
+        mController.onNativeInitialized();
+        verify(mTask, times(2)).finish(anyObject(), anyObject());
+    }
+
+    @Test
+    public void testReparenting_twoTabsOneIncognito() {
+        createAndAddMockTab(1, false);
+        mForegroundTab = createAndAddMockTab(2, true);
+        mController.onNightModeStateChanged();
+
+        AsyncTabParams params = AsyncTabParamsManager.getAsyncTabParams().get(2);
+        Assert.assertNotNull(params);
+        Assert.assertTrue(params instanceof TabReparentingParams);
+
+        TabReparentingParams trp = (TabReparentingParams) params;
+        Assert.assertEquals(
+                "The index of the first tab stored should match its index in the tab stack.", 0,
+                trp.getTabIndex());
+        Assert.assertTrue(trp.isFromNightModeReparenting());
+        Assert.assertTrue(trp.isForegroundTab());
+
+        Tab tab = trp.getTabToReparent();
+        Assert.assertNotNull(tab);
+        Assert.assertEquals(2, tab.getId());
 
         verify(mTask, times(2)).detach();
 
@@ -216,14 +227,14 @@ public class NightModeReparentingControllerTest {
 
     @Test
     public void testReparenting_threeTabsOutOfOrder() {
-        createAndAddMockTab(1, 1);
-        createAndAddMockTab(2, 0);
-        createAndAddMockTab(3, 2);
+        createAndAddMockTab(3, false);
+        mForegroundTab = createAndAddMockTab(2, false);
+        createAndAddMockTab(1, false);
         mController.onNightModeStateChanged();
 
         // Check the foreground tab.
         TabReparentingParams trp =
-                (TabReparentingParams) AsyncTabParamsManager.getAsyncTabParams().get(1);
+                (TabReparentingParams) AsyncTabParamsManager.getAsyncTabParams().get(2);
         Assert.assertEquals(
                 "The index of the first tab stored should match its index in the tab stack.", 1,
                 trp.getTabIndex());
@@ -232,10 +243,10 @@ public class NightModeReparentingControllerTest {
 
         Tab tab = trp.getTabToReparent();
         Assert.assertNotNull(tab);
-        Assert.assertEquals(1, tab.getId());
+        Assert.assertEquals(2, tab.getId());
 
         // Check the background tabs.
-        trp = (TabReparentingParams) AsyncTabParamsManager.getAsyncTabParams().get(2);
+        trp = (TabReparentingParams) AsyncTabParamsManager.getAsyncTabParams().get(1);
         Assert.assertFalse("The index of the background tabs stored shouldn't have a tab index.",
                 trp.hasTabIndex());
         Assert.assertTrue(trp.isFromNightModeReparenting());
@@ -265,30 +276,34 @@ public class NightModeReparentingControllerTest {
         }
     }
 
-    private void createAndAddMockTab(int id, int index) {
+    /**
+     * Adds a tab to the correct model and sets the index in the mapping.
+     *
+     * @param id The id to give to the tab.
+     * @param incognito Whether to add the tab to the incognito model or the regular model.
+     * @return The tab that was added. Use the return value to set the foreground tab for tests.
+     */
+    private Tab createAndAddMockTab(int id, boolean incognito) {
         Tab tab = Mockito.mock(Tab.class);
         UserDataHost udh = new UserDataHost();
         udh.setUserData(ReparentingTask.class, mTask);
         doReturn(udh).when(tab).getUserDataHost();
         doReturn(id).when(tab).getId();
-        mTabs.add(tab);
-        mTabIndexMapping.put(tab, index);
-    }
 
-    private Tab getNextTabFromList() {
-        if (mTabs.size() == 0) return null;
-        return mTabs.get(0);
-    }
-
-    private void removeTab(Tab tab) {
-        mTabs.remove(tab);
-    }
+        int index;
+        if (incognito) {
+            mIncognitoTabModel.addTab(tab, -1, TabLaunchType.FROM_BROWSER_ACTIONS);
+            index = mIncognitoTabModel.indexOf(tab);
+        } else {
+            mTabModel.addTab(tab, -1, TabLaunchType.FROM_BROWSER_ACTIONS);
+            index = mTabModel.indexOf(tab);
+        }
+        mTabIndexMapping.put(tab, index);
 
-    private TabModel getCurrentTabModel() {
-        return mTabModel;
+        return tab;
     }
 
-    private int getTabIndex(Tab tab) {
-        return mTabIndexMapping.get(tab);
+    private Tab getForegroundTab() {
+        return mForegroundTab;
     }
 }