From 49d4429881772e353ccadda641d3f78b25cb3fec Mon Sep 17 00:00:00 2001
From: Jagrut Desai <jagrutdesai@google.com>
Date: Tue, 9 Apr 2024 12:49:15 -0700
Subject: [PATCH] Fixing TaplUninstallRemoveTest#uninstallWorkspaceIcon test

Test: Presubmit
Bug: 319501259
Flag: NONE
Change-Id: I9566b947bc7d5d64953a8c637fefd4fa37cca87c
---
 .../launcher3/dragging/TaplUninstallRemoveTest.java  | 12 ++++++------
 tests/tapl/com/android/launcher3/tapl/Workspace.java |  1 +
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
index df71ec09d4..b7933c8aff 100644
--- a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
+++ b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
@@ -39,6 +39,7 @@ import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape;
 import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.Wait;
+import com.android.launcher3.util.rule.ScreenRecordRule;
 import com.android.launcher3.util.rule.TestStabilityRule;
 
 import org.junit.Test;
@@ -132,6 +133,7 @@ public class TaplUninstallRemoveTest extends AbstractLauncherUiTest<Launcher> {
      */
     @Test
     @PlatinumTest(focusArea = "launcher")
+    @ScreenRecordRule.ScreenRecord // b/319501259
     public void uninstallWorkspaceIcon() throws IOException {
         Point[] gridPositions = TestUtil.getCornersAndCenterPositions(mLauncher);
         StringBuilder sb = new StringBuilder();
@@ -162,12 +164,10 @@ public class TaplUninstallRemoveTest extends AbstractLauncherUiTest<Launcher> {
             mLauncher.getWorkspace().verifyWorkspaceAppIconIsGone(
                     DUMMY_APP_NAME + " was expected to disappear after uninstall.", DUMMY_APP_NAME);
 
-            if (!TestStabilityRule.isPresubmit()) { // b/315847371
-                Log.d(UIOBJECT_STALE_ELEMENT, "second getWorkspaceIconsPositions()");
-                Map<String, Point> finalPositions =
-                        mLauncher.getWorkspace().getWorkspaceIconsPositions();
-                assertThat(finalPositions).doesNotContainKey(DUMMY_APP_NAME);
-            }
+            Log.d(UIOBJECT_STALE_ELEMENT, "second getWorkspaceIconsPositions()");
+            Map<String, Point> finalPositions =
+                    mLauncher.getWorkspace().getWorkspaceIconsPositions();
+            assertThat(finalPositions).doesNotContainKey(DUMMY_APP_NAME);
         } finally {
             TestUtil.uninstallDummyApp();
         }
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index d176136e19..4fa93ef2ca 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -347,6 +347,7 @@ public final class Workspace extends Home {
      */
     public Map<String, Point> getWorkspaceIconsPositions() {
         final UiObject2 workspace = verifyActiveContainer();
+        mLauncher.waitForLauncherInitialized(); // b/319501259
         List<UiObject2> workspaceIcons =
                 mLauncher.waitForObjectsInContainer(workspace, AppIcon.getAnyAppIconSelector());
         return workspaceIcons.stream()