Add resume check to other WorkProfileTests to fix flakiness.

- Tests were flaky with waitForStateTransitionToEnd
- Tests stopped flaking with both waitForResume and waitForStateTransitionToEnd
- waitForResume may be enough to fix the flakiness, so
  removing waitForStateTransitionToEnd for now, if flakiness
  returns we can add it back

Bug: 202735477
Test: presubmit
Change-Id: Ib9cc16ea7815387582486517a87249d0053475f6
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 44f2719..314aee4 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -352,14 +352,6 @@
                 launcher -> launcher.getStateManager().getCurrentStableState() == state.get());
     }
 
-    // Cannot be used in TaplTests between a Tapl call injecting a gesture and a tapl call
-    // expecting the results of that gesture because the wait can hide flakeness.
-    protected void waitForStateTransitionToEnd(String message, Supplier<LauncherState> state) {
-        waitForLauncherCondition(message,
-                launcher -> launcher.getStateManager().isInStableState(state.get())
-                        && !launcher.getStateManager().isInTransition());
-    }
-
     protected void waitForResumed(String message) {
         waitForLauncherCondition(message, launcher -> launcher.hasBeenResumed());
     }
diff --git a/tests/src/com/android/launcher3/ui/WorkProfileTest.java b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
index 3a338d1..4bb6e91 100644
--- a/tests/src/com/android/launcher3/ui/WorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
@@ -59,6 +59,13 @@
         String[] tokens = output.split("\\s+");
         mProfileUserId = Integer.parseInt(tokens[tokens.length - 1]);
         mDevice.executeShellCommand("am start-user " + mProfileUserId);
+
+        mDevice.pressHome();
+        waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
+        waitForState("Launcher internal state didn't switch to Normal", () -> NORMAL);
+        waitForResumed("Launcher internal state is still Background");
+        executeOnLauncher(launcher -> launcher.getStateManager().goToState(ALL_APPS));
+        waitForState("Launcher internal state didn't switch to All Apps", () -> ALL_APPS);
     }
 
     @After
@@ -89,14 +96,6 @@
     @Test
     @ScreenRecord // b/202735477
     public void workTabExists() {
-        mDevice.pressHome();
-        waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
-        waitForStateTransitionToEnd("Launcher internal state didn't switch to Normal",
-                () -> NORMAL);
-        waitForResumed("Launcher internal state is still Background");
-        executeOnLauncher(launcher -> launcher.getStateManager().goToState(ALL_APPS));
-        waitForStateTransitionToEnd("Launcher internal state didn't switch to All Apps",
-                () -> ALL_APPS);
         waitForLauncherCondition("Personal tab is missing",
                 launcher -> launcher.getAppsView().isPersonalTabVisible(),
                 LauncherInstrumentation.WAIT_TIME_MS);
@@ -107,12 +106,6 @@
 
     @Test
     public void toggleWorks() {
-        mDevice.pressHome();
-        waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
-        waitForState("Launcher internal state didn't switch to Normal", () -> NORMAL);
-        executeOnLauncher(launcher -> launcher.getStateManager().goToState(ALL_APPS));
-        waitForState("Launcher internal state didn't switch to All Apps", () -> ALL_APPS);
-
         waitForWorkTabSetup();
 
         executeOnLauncher(launcher -> {
@@ -154,11 +147,6 @@
 
     @Test
     public void testEdu() {
-        mDevice.pressHome();
-        waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
-        waitForState("Launcher internal state didn't switch to Normal", () -> NORMAL);
-        executeOnLauncher(launcher -> launcher.getStateManager().goToState(ALL_APPS));
-        waitForState("Launcher internal state didn't switch to All Apps", () -> ALL_APPS);
         waitForWorkTabSetup();
         executeOnLauncher(l -> {
             l.getSharedPrefs().edit().putInt(WorkAdapterProvider.KEY_WORK_EDU_STEP, 0).commit();