Allow workspace to provide icon positions in tests

This is needed in order to implement a test for the unfold animation.
getWorkspaceAndHotseatIconsPositions returns map that associate icon name with its center on the screen

Test: atest TaplTestsLauncher3
Bug: 220087235
Change-Id: I254aa0c146edf4cae1c2bf5c09dec000993eb4d8
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 7080c85..6f8b9d2 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -613,6 +613,10 @@
         return createShortcutIfNotExist(name, dimension.x / 2, dimension.y / 2);
     }
 
+    protected HomeAppIcon createShortcutIfNotExist(String name, Point cellPosition) {
+        return createShortcutIfNotExist(name, cellPosition.x, cellPosition.y);
+    }
+
     protected HomeAppIcon createShortcutIfNotExist(String name, int cellX, int cellY) {
         HomeAppIcon homeAppIcon = mLauncher.getWorkspace().tryGetWorkspaceAppIcon(name);
         if (homeAppIcon == null) {
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index f0bef24..15e8f68 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -18,6 +18,8 @@
 
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -54,11 +56,16 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.io.IOException;
+import java.util.Map;
+
 @LargeTest
 @RunWith(AndroidJUnit4.class)
 public class TaplTestsLauncher3 extends AbstractLauncherUiTest {
     private static final String APP_NAME = "LauncherTestApp";
     private static final String DUMMY_APP_NAME = "Aardwolf";
+    private static final String MAPS_APP_NAME = "Maps";
+    private static final String STORE_APP_NAME = "Play Store";
 
     @Before
     public void setUp() throws Exception {
@@ -462,15 +469,7 @@
     @Test
     @PortraitLandscape
     public void testDragAppIconToWorkspaceCell() throws Exception {
-        final Point dimensions = mLauncher.getWorkspace().getIconGridDimensions();
-
-        Point[] targets = {
-                new Point(0, 1),
-                new Point(0, dimensions.y - 2),
-                new Point(dimensions.x - 1, 1),
-                new Point(dimensions.x - 1, dimensions.y - 2),
-                new Point(dimensions.x / 2, dimensions.y / 2)
-        };
+        Point[] targets = getCornersAndCenterPositions();
 
         for (Point target : targets) {
             final HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
@@ -491,6 +490,48 @@
         }
     }
 
+    @Test
+    public void getIconsPosition_afterIconRemoved_notContained() throws IOException {
+        Point[] gridPositions = getCornersAndCenterPositions();
+        createShortcutIfNotExist(STORE_APP_NAME, gridPositions[0]);
+        createShortcutIfNotExist(MAPS_APP_NAME, gridPositions[1]);
+        TestUtil.installDummyApp();
+        try {
+            createShortcutIfNotExist(DUMMY_APP_NAME, gridPositions[2]);
+            Map<String, Point> initialPositions =
+                    mLauncher.getWorkspace().getWorkspaceIconsPositions();
+            assertThat(initialPositions.keySet())
+                    .containsAtLeast(DUMMY_APP_NAME, MAPS_APP_NAME, STORE_APP_NAME);
+
+            mLauncher.getWorkspace().getWorkspaceAppIcon(DUMMY_APP_NAME).uninstall();
+
+            assertNull(
+                    DUMMY_APP_NAME + " app was found after being uninstalled",
+                    mLauncher.getWorkspace().tryGetWorkspaceAppIcon(DUMMY_APP_NAME));
+
+            Map<String, Point> finalPositions =
+                    mLauncher.getWorkspace().getWorkspaceIconsPositions();
+            assertThat(finalPositions).doesNotContainKey(DUMMY_APP_NAME);
+        } finally {
+            TestUtil.uninstallDummyApp();
+        }
+    }
+
+    /**
+     * @return List of workspace grid coordinates. Those are not pixels. See {@link
+     *     Workspace#getIconGridDimensions()}
+     */
+    private Point[] getCornersAndCenterPositions() {
+        final Point dimensions = mLauncher.getWorkspace().getIconGridDimensions();
+        return new Point[] {
+            new Point(0, 1),
+            new Point(0, dimensions.y - 2),
+            new Point(dimensions.x - 1, 1),
+            new Point(dimensions.x - 1, dimensions.y - 2),
+            new Point(dimensions.x / 2, dimensions.y / 2)
+        };
+    }
+
     public static String getAppPackageName() {
         return getInstrumentation().getContext().getPackageName();
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index ae0f6a6..b0b42ba 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -43,6 +43,7 @@
 import com.android.launcher3.testing.WorkspaceCellCenterRequest;
 
 import java.util.List;
+import java.util.Map;
 import java.util.function.Supplier;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -222,6 +223,21 @@
                 mHotseat, AppIcon.getAppIconSelector(appName, mLauncher)));
     }
 
+    /**
+     * @return map of text -> center of the view. In case of icons with the same name, the one with
+     *     lower x coordinate is selected.
+     */
+    public Map<String, Point> getWorkspaceIconsPositions() {
+        final UiObject2 workspace = verifyActiveContainer();
+        List<UiObject2> workspaceIcons =
+                mLauncher.waitForObjectsInContainer(workspace, AppIcon.getAnyAppIconSelector());
+        return workspaceIcons.stream()
+                .collect(
+                        Collectors.toMap(
+                                /* keyMapper= */ UiObject2::getText,
+                                /* valueMapper= */ UiObject2::getVisibleCenter,
+                                /* mergeFunction= */ (p1, p2) -> p1.x < p2.x ? p1 : p2));
+    }
     /*
      * Get the center point of the delete/uninstall icon in the drop target bar.
      */