Merge "Making ReorderWidget test more reliable" into tm-qpr-dev
diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index d31a646..0b756b6 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -208,6 +208,11 @@
                 );
             }
 
+            case TestProtocol.REQUEST_WORKSPACE_CURRENT_PAGE_INDEX: {
+                return getLauncherUIProperty(Bundle::putInt,
+                        launcher -> launcher.getWorkspace().getCurrentPage());
+            }
+
             case TestProtocol.REQUEST_HOTSEAT_CELL_CENTER: {
                 final HotseatCellCenterRequest request = extra.getParcelable(
                         TestProtocol.TEST_INFO_REQUEST_FIELD);
diff --git a/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java b/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java
index 11363a2..65873f1 100644
--- a/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java
+++ b/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java
@@ -126,6 +126,9 @@
     public static final String REQUEST_WORKSPACE_CELL_CENTER = "workspace-cell-center";
     public static final String REQUEST_WORKSPACE_COLUMNS_ROWS = "workspace-columns-rows";
 
+    public static final String REQUEST_WORKSPACE_CURRENT_PAGE_INDEX =
+            "workspace-current-page-index";
+
     public static final String REQUEST_HOTSEAT_CELL_CENTER = "hotseat-cell-center";
 
     public static final String REQUEST_GET_FOCUSED_TASK_HEIGHT_FOR_TABLET =
diff --git a/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java b/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java
index 3443fd9..77fca96 100644
--- a/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java
+++ b/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java
@@ -114,6 +114,10 @@
                 new FavoriteItemsTransaction(mTargetContext, this);
         transaction = buildWorkspaceFromBoards(testCase.mStart, transaction);
         transaction.commit();
+        // resetLoaderState triggers the launcher to start loading the workspace which allows
+        // waitForLauncherCondition to wait for that condition, otherwise the condition would
+        // always be true and it wouldn't wait for the changes to be applied.
+        resetLoaderState();
         waitForLauncherCondition("Workspace didn't finish loading", l -> !l.isWorkspaceLoading());
         Widget widget = mLauncher.getWorkspace().getWidgetAtCell(mainWidgetCellPos.getCellX(),
                 mainWidgetCellPos.getCellY());
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 473cfb9..0197a11 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -67,6 +67,7 @@
             "Key event: KeyEvent.*?action=ACTION_UP.*?keyCode=KEYCODE_W"
                     + ".*?metaState=META_CTRL_ON");
     static final Pattern LONG_CLICK_EVENT = Pattern.compile("onWorkspaceItemLongClick");
+    public static final int MAX_WORKSPACE_DRAG_TRIES = 100;
 
     private final UiObject2 mHotseat;
 
@@ -430,6 +431,12 @@
                 TestProtocol.TEST_INFO_RESPONSE_FIELD);
     }
 
+    /** Returns the index of the current page */
+    static int geCurrentPage(LauncherInstrumentation launcher) {
+        return launcher.getTestInfo(TestProtocol.REQUEST_WORKSPACE_CURRENT_PAGE_INDEX).getInt(
+                TestProtocol.TEST_INFO_RESPONSE_FIELD);
+    }
+
     /**
      * Finds folder icons in the current workspace.
      *
@@ -569,21 +576,10 @@
                     expectLongClickEvents,
                     /* runToSpringLoadedState= */ true);
             Point targetDest = getCellCenter(launcher, cellX, cellY, spanX, spanY);
-            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
-            for (int i = 0; i < destinationWorkspace; i++) {
-                // Don't drag all the way to the edge to prevent touch events from getting out of
-                //screen bounds.
-                Point screenEdge = new Point(displayX - 1, targetDest.y);
-                Point finalDragStart = dragStart;
-                executeAndWaitForPageScroll(launcher,
-                        () -> launcher.movePointer(finalDragStart, screenEdge, DEFAULT_DRAG_STEPS,
-                                true, downTime, downTime, true,
-                                LauncherInstrumentation.GestureScope.INSIDE));
-                dragStart = screenEdge;
-            }
+            dragStart = dragToGivenWorkspace(launcher, dragStart, destinationWorkspace,
+                    targetDest.y);
 
             // 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
@@ -595,6 +591,45 @@
         }
     }
 
+    /**
+     * Given a drag that already started at currentPosition, drag the item to the given destination
+     * index defined by destinationWorkspaceIndex.
+     *
+     * @param launcher
+     * @param currentPosition
+     * @param destinationWorkspaceIndex
+     * @param y
+     * @return the finishing position of the drag.
+     */
+    static Point dragToGivenWorkspace(LauncherInstrumentation launcher, Point currentPosition,
+            int destinationWorkspaceIndex, int y) {
+        final long downTime = SystemClock.uptimeMillis();
+        int displayX = launcher.getRealDisplaySize().x;
+        int currentPage = Workspace.geCurrentPage(launcher);
+        int counter = 0;
+        while (currentPage != destinationWorkspaceIndex) {
+            counter++;
+            if (counter > MAX_WORKSPACE_DRAG_TRIES) {
+                throw new RuntimeException(
+                        "Wrong destination workspace index " + destinationWorkspaceIndex
+                                + ", desired workspace was never reached");
+            }
+            // if the destination is greater than current page, set the display edge to be the
+            // right edge. Don't drag all the way to the edge to prevent touch events from
+            // getting out of screen bounds.
+            int displayEdge = destinationWorkspaceIndex > currentPage ? displayX - 1 : 1;
+            Point screenEdge = new Point(displayEdge, y);
+            Point finalDragStart = currentPosition;
+            executeAndWaitForPageScroll(launcher,
+                    () -> launcher.movePointer(finalDragStart, screenEdge, DEFAULT_DRAG_STEPS,
+                            true, downTime, downTime, true,
+                            LauncherInstrumentation.GestureScope.INSIDE));
+            currentPage = Workspace.geCurrentPage(launcher);
+            currentPosition = screenEdge;
+        }
+        return currentPosition;
+    }
+
     private static void executeAndWaitForPageScroll(LauncherInstrumentation launcher,
             Runnable command) {
         launcher.executeAndWaitForEvent(command,