Add TAPL tests for two panel workspace

- Add test for icon dragging
- Add test for page pairing
- Add test for empty page removal
- Add dragIcon helper method

Test: TwoPanelWorkspaceTest
Bug: 197631877
Change-Id: I7ce2dd88e183626feac49b34254d00a589473944
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index c90d283..41c7c37 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -80,12 +80,8 @@
         assertTrue(message, failed);
     }
 
-    private int pagesPerScreen() {
-        return mLauncher.isTwoPanels() ? 2 : 1;
-    }
-
-    private boolean isWorkspaceScrollable(Launcher launcher) {
-        return launcher.getWorkspace().getPageCount() > pagesPerScreen();
+    public static boolean isWorkspaceScrollable(Launcher launcher) {
+        return launcher.getWorkspace().getPageCount() > launcher.getWorkspace().getPanelCount();
     }
 
     private int getCurrentWorkspacePage(Launcher launcher) {
@@ -192,7 +188,7 @@
         executeOnLauncher(
                 launcher -> assertEquals(
                         "Ensuring workspace scrollable didn't switch to next screen",
-                        pagesPerScreen(), getCurrentWorkspacePage(launcher)));
+                        workspace.pagesPerScreen(), getCurrentWorkspacePage(launcher)));
         executeOnLauncher(
                 launcher -> assertTrue("ensureScrollable didn't make workspace scrollable",
                         isWorkspaceScrollable(launcher)));
@@ -209,7 +205,7 @@
         workspace.flingForward();
         executeOnLauncher(
                 launcher -> assertEquals("Flinging forward didn't switch workspace to next screen",
-                        pagesPerScreen(), getCurrentWorkspacePage(launcher)));
+                        workspace.pagesPerScreen(), getCurrentWorkspacePage(launcher)));
         assertTrue("Launcher internal state is not Home", isInState(() -> LauncherState.NORMAL));
 
         // Test starting a workspace app.
diff --git a/tests/src/com/android/launcher3/ui/workspace/TwoPanelWorkspaceTest.java b/tests/src/com/android/launcher3/ui/workspace/TwoPanelWorkspaceTest.java
new file mode 100644
index 0000000..b048cd4
--- /dev/null
+++ b/tests/src/com/android/launcher3/ui/workspace/TwoPanelWorkspaceTest.java
@@ -0,0 +1,285 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.ui.workspace;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import androidx.test.filters.LargeTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.launcher3.CellLayout;
+import com.android.launcher3.Launcher;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.tapl.Workspace;
+import com.android.launcher3.ui.AbstractLauncherUiTest;
+import com.android.launcher3.ui.TaplTestsLauncher3;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Arrays;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Tests for two panel workspace.
+ *
+ * Note running these tests will clear the workspace on the device.
+ */
+@LargeTest
+@RunWith(AndroidJUnit4.class)
+public class TwoPanelWorkspaceTest extends AbstractLauncherUiTest {
+
+    Workspace mWorkspace;
+
+    @Before
+    public void setUp() throws Exception {
+        super.setUp();
+        TaplTestsLauncher3.initialize(this);
+        mWorkspace = mLauncher.getWorkspace();
+    }
+
+    @Test
+    public void testDragIconToRightPanel() {
+        if (!mLauncher.isTwoPanels()) {
+            return;
+        }
+
+        // Pre verifying the screens
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1);
+            assertItemsOnPage(launcher, 0, "Play Store", "Maps");
+            assertPageEmpty(launcher, 1);
+        });
+
+        mWorkspace.dragIcon(mWorkspace.getHotseatAppIcon("Chrome"), 1);
+
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1);
+            assertItemsOnPage(launcher, 0, "Maps", "Play Store");
+            assertItemsOnPage(launcher, 1, "Chrome");
+        });
+    }
+
+    @Test
+    public void testDragIconToPage2() {
+        if (!mLauncher.isTwoPanels()) {
+            return;
+        }
+
+        // Pre verifying the screens
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1);
+            assertItemsOnPage(launcher, 0, "Play Store", "Maps");
+            assertPageEmpty(launcher, 1);
+        });
+
+        mWorkspace.dragIcon(mWorkspace.getWorkspaceAppIcon("Maps"), 2);
+
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1, 2, 3);
+            assertItemsOnPage(launcher, 0, "Play Store");
+            assertPageEmpty(launcher, 1);
+            assertItemsOnPage(launcher, 2, "Maps");
+            assertPageEmpty(launcher, 3);
+        });
+    }
+
+    @Test
+    public void testDragIconToPage3() {
+        if (!mLauncher.isTwoPanels()) {
+            return;
+        }
+
+        // Pre verifying the screens
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1);
+            assertItemsOnPage(launcher, 0, "Play Store", "Maps");
+            assertPageEmpty(launcher, 1);
+        });
+
+        mWorkspace.dragIcon(mWorkspace.getHotseatAppIcon("Phone"), 3);
+
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1, 2, 3);
+            assertItemsOnPage(launcher, 0, "Play Store", "Maps");
+            assertPageEmpty(launcher, 1);
+            assertPageEmpty(launcher, 2);
+            assertItemsOnPage(launcher, 3, "Phone");
+        });
+    }
+
+
+    @Test
+    public void testEmptyPageDoesNotGetRemovedIfPagePairIsNotEmpty() {
+        if (!mLauncher.isTwoPanels()) {
+            return;
+        }
+
+        // Pre verifying the screens
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1);
+            assertItemsOnPage(launcher, 0, "Play Store", "Maps");
+            assertPageEmpty(launcher, 1);
+        });
+
+        mWorkspace.dragIcon(mWorkspace.getWorkspaceAppIcon("Maps"), 3);
+        mWorkspace.dragIcon(mWorkspace.getHotseatAppIcon("Chrome"), 0);
+
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1, 2, 3);
+            assertItemsOnPage(launcher, 0, "Play Store");
+            assertPageEmpty(launcher, 1);
+            assertItemsOnPage(launcher, 2, "Chrome");
+            assertItemsOnPage(launcher, 3, "Maps");
+        });
+
+        mWorkspace.dragIcon(mWorkspace.getWorkspaceAppIcon("Maps"), -1);
+
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1, 2, 3);
+            assertItemsOnPage(launcher, 0, "Play Store");
+            assertItemsOnPage(launcher, 1, "Maps");
+            assertItemsOnPage(launcher, 2, "Chrome");
+            assertPageEmpty(launcher, 3);
+        });
+
+        // Move Chrome to the right panel as well, to make sure pages are not deleted whichever
+        // page is the empty one
+        mWorkspace.flingForward();
+        mWorkspace.dragIcon(mWorkspace.getWorkspaceAppIcon("Chrome"), 1);
+
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1, 2, 3);
+            assertItemsOnPage(launcher, 0, "Play Store");
+            assertItemsOnPage(launcher, 1, "Maps");
+            assertPageEmpty(launcher, 2);
+            assertItemsOnPage(launcher, 3, "Chrome");
+        });
+    }
+
+
+    @Test
+    public void testEmptyPagesGetRemovedIfBothPagesAreEmpty() {
+        if (!mLauncher.isTwoPanels()) {
+            return;
+        }
+
+        // Pre verifying the screens
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1);
+            assertItemsOnPage(launcher, 0, "Play Store", "Maps");
+            assertPageEmpty(launcher, 1);
+        });
+
+        mWorkspace.dragIcon(mWorkspace.getWorkspaceAppIcon("Play Store"), 2);
+        mWorkspace.dragIcon(mWorkspace.getHotseatAppIcon("Camera"), 1);
+
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1, 2, 3);
+            assertItemsOnPage(launcher, 0, "Maps");
+            assertPageEmpty(launcher, 1);
+            assertItemsOnPage(launcher, 2, "Play Store");
+            assertItemsOnPage(launcher, 3, "Camera");
+        });
+
+        mWorkspace.dragIcon(mWorkspace.getWorkspaceAppIcon("Camera"), -1);
+        mWorkspace.flingForward();
+        mWorkspace.dragIcon(mWorkspace.getWorkspaceAppIcon("Play Store"), -2);
+
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1);
+            assertItemsOnPage(launcher, 0, "Play Store", "Maps");
+            assertItemsOnPage(launcher, 1, "Camera");
+        });
+    }
+
+    @Test
+    public void testMiddleEmptyPagesGetRemoved() {
+        if (!mLauncher.isTwoPanels()) {
+            return;
+        }
+
+        // Pre verifying the screens
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1);
+            assertItemsOnPage(launcher, 0, "Play Store", "Maps");
+            assertPageEmpty(launcher, 1);
+        });
+
+        mWorkspace.dragIcon(mWorkspace.getWorkspaceAppIcon("Maps"), 2);
+        mWorkspace.dragIcon(mWorkspace.getHotseatAppIcon("Messages"), 3);
+
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1, 2, 3, 4, 5);
+            assertItemsOnPage(launcher, 0, "Play Store");
+            assertPageEmpty(launcher, 1);
+            assertItemsOnPage(launcher, 2, "Maps");
+            assertPageEmpty(launcher, 3);
+            assertPageEmpty(launcher, 4);
+            assertItemsOnPage(launcher, 5, "Messages");
+        });
+
+        mWorkspace.flingBackward();
+        mWorkspace.dragIcon(mWorkspace.getWorkspaceAppIcon("Maps"), 2);
+
+        executeOnLauncher(launcher -> {
+            assertPagesExist(launcher, 0, 1, 4, 5);
+            assertItemsOnPage(launcher, 0, "Play Store");
+            assertPageEmpty(launcher, 1);
+            assertItemsOnPage(launcher, 4, "Maps");
+            assertItemsOnPage(launcher, 5, "Messages");
+        });
+    }
+
+    private void assertPageEmpty(Launcher launcher, int pageId) {
+        CellLayout page = launcher.getWorkspace().getScreenWithId(pageId);
+        assertNotNull("Page " + pageId + " does NOT exist.", page);
+        assertEquals("Page " + pageId + " is NOT empty. Number of items on the page:", 0,
+                page.getShortcutsAndWidgets().getChildCount());
+    }
+
+    private void assertPagesExist(Launcher launcher, int... pageIds) {
+        int pageCount = launcher.getWorkspace().getPageCount();
+        assertEquals("Existing page count does NOT match.", pageIds.length, pageCount);
+        for (int i = 0; i < pageCount; i++) {
+            CellLayout page = (CellLayout) launcher.getWorkspace().getPageAt(i);
+            int pageId = launcher.getWorkspace().getIdForScreen(page);
+            assertEquals("The page's id at index " + i + " does NOT match.", pageId,
+                    pageIds[i]);
+        }
+    }
+
+    private void assertItemsOnPage(Launcher launcher, int pageId, String... itemTitles) {
+        Set<String> itemTitleSet = Arrays.stream(itemTitles).collect(Collectors.toSet());
+        CellLayout page = launcher.getWorkspace().getScreenWithId(pageId);
+        int itemCount = page.getShortcutsAndWidgets().getChildCount();
+        for (int i = 0; i < itemCount; i++) {
+            ItemInfo itemInfo = (ItemInfo) page.getShortcutsAndWidgets().getChildAt(i).getTag();
+            if (itemInfo != null) {
+                assertTrue("There was an extra item on page " + pageId + ": " + itemInfo.title,
+                        itemTitleSet.remove(itemInfo.title));
+            }
+        }
+        assertTrue("Could NOT find some of the items on page " + pageId + ": "
+                        + itemTitleSet.stream().collect(Collectors.joining(",")),
+                itemTitleSet.isEmpty());
+    }
+}
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 0145690..d9f5cc8 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -145,16 +145,7 @@
             if (!isWorkspaceScrollable(workspace)) {
                 try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                         "dragging icon to a second page of workspace to make it scrollable")) {
-                    dragIconToWorkspace(
-                            mLauncher,
-                            getHotseatAppIcon("Chrome"),
-                            new Point(mLauncher.getDevice().getDisplayWidth(),
-                                    mLauncher.getVisibleBounds(workspace).centerY()),
-                            "popup_container",
-                            false,
-                            false,
-                            () -> mLauncher.expectEvent(
-                                    TestProtocol.SEQUENCE_MAIN, LONG_CLICK_EVENT));
+                    dragIcon(workspace, getHotseatAppIcon("Chrome"), pagesPerScreen());
                     verifyActiveContainer();
                 }
             }
@@ -163,6 +154,48 @@
         }
     }
 
+    /**
+     * Returns the number of pages that are visible on the screen simultaneously.
+     */
+    public int pagesPerScreen() {
+        return mLauncher.isTwoPanels() ? 2 : 1;
+    }
+
+    /**
+     * Drags an icon to the (currentPage + pageDelta) page if the page already exists.
+     * If the target page doesn't exist, the icon will be put onto an existing page that is the
+     * closest to the target page.
+     *
+     * @param appIcon   - icon to drag.
+     * @param pageDelta - how many pages should the icon be dragged from the current page.
+     *                    It can be a negative value.
+     */
+    public void dragIcon(AppIcon appIcon, int pageDelta) {
+        try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
+            final UiObject2 workspace = verifyActiveContainer();
+            try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                    "dragging icon to page with delta: " + pageDelta)) {
+                dragIcon(workspace, appIcon, pageDelta);
+                verifyActiveContainer();
+            }
+        }
+    }
+
+    private void dragIcon(UiObject2 workspace, AppIcon appIcon, int pageDelta) {
+        int pageWidth = mLauncher.getDevice().getDisplayWidth() / pagesPerScreen();
+        int targetX = (pageWidth / 2) + pageWidth * pageDelta;
+        dragIconToWorkspace(
+                mLauncher,
+                appIcon,
+                new Point(targetX, mLauncher.getVisibleBounds(workspace).centerY()),
+                "popup_container",
+                false,
+                false,
+                () -> mLauncher.expectEvent(
+                        TestProtocol.SEQUENCE_MAIN, LONG_CLICK_EVENT));
+        verifyActiveContainer();
+    }
+
     private boolean isWorkspaceScrollable(UiObject2 workspace) {
         return workspace.getChildCount() > (mLauncher.isTwoPanels() ? 2 : 1);
     }
@@ -258,10 +291,26 @@
         try (LauncherInstrumentation.Closable ignored = launcher.addContextLayer(
                 "want to drag icon to workspace")) {
             final long downTime = SystemClock.uptimeMillis();
-            final Point dragStartCenter = dragIconToSpringLoaded(launcher, downTime,
+            Point dragStart = dragIconToSpringLoaded(launcher, downTime,
                     launchable.getObject(), longPressIndicator, expectLongClickEvents);
-            final Point targetDest = dest.get();
-            launcher.movePointer(dragStartCenter, targetDest, 10, downTime, true,
+            Point targetDest = dest.get();
+            int displayX = launcher.getRealDisplaySize().x;
+
+            // Since the destination can be on another page, we need to drag to the edge first
+            // until we reach the target page
+            while (targetDest.x > displayX || targetDest.x < 0) {
+                int edgeX = targetDest.x > 0 ? displayX : 0;
+                Point screenEdge = new Point(edgeX, targetDest.y);
+                launcher.movePointer(dragStart, screenEdge, 10, downTime, true,
+                        LauncherInstrumentation.GestureScope.INSIDE);
+                launcher.waitForIdle(); // Wait for the page change to happen
+                targetDest.x += displayX * (targetDest.x > 0 ? -1 : 1);
+                dragStart = screenEdge;
+            }
+
+            // targetDest.x is now between 0 and displayX so we found the target page,
+            // we just have to put move the icon to the destination and drop it
+            launcher.movePointer(dragStart, targetDest, 10, downTime, true,
                     LauncherInstrumentation.GestureScope.INSIDE);
             dropDraggedIcon(launcher, targetDest, downTime, expectDropEvents);
         }