diff --git a/src/com/android/launcher3/statemanager/StateManager.java b/src/com/android/launcher3/statemanager/StateManager.java
index b34af975bd..1767939114 100644
--- a/src/com/android/launcher3/statemanager/StateManager.java
+++ b/src/com/android/launcher3/statemanager/StateManager.java
@@ -84,7 +84,7 @@ public class StateManager<STATE_TYPE extends BaseState<STATE_TYPE>> {
                 + ", mCurrentStableState:" + mCurrentStableState
                 + ", mState:" + mState
                 + ", mRestState:" + mRestState
-                + ", isInTransition:" + (mConfig.currentAnimation != null) + ")";
+                + ", isInTransition:" + isInTransition() + ")";
     }
 
     public void dump(String prefix, PrintWriter writer) {
@@ -93,7 +93,7 @@ public class StateManager<STATE_TYPE extends BaseState<STATE_TYPE>> {
         writer.println(prefix + "\tmCurrentStableState:" + mCurrentStableState);
         writer.println(prefix + "\tmState:" + mState);
         writer.println(prefix + "\tmRestState:" + mRestState);
-        writer.println(prefix + "\tisInTransition:" + (mConfig.currentAnimation != null));
+        writer.println(prefix + "\tisInTransition:" + isInTransition());
     }
 
     public StateHandler[] getStateHandlers() {
@@ -129,6 +129,13 @@ public class StateManager<STATE_TYPE extends BaseState<STATE_TYPE>> {
                 && (mConfig.targetState == null || mConfig.targetState == state);
     }
 
+    /**
+     * @return {@code true} If there is an active transition.
+     */
+    public boolean isInTransition() {
+        return mConfig.currentAnimation != null;
+    }
+
     /**
      * @see #goToState(STATE_TYPE, boolean, AnimatorListener)
      */
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index b3457cda75..44f2719a02 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -354,9 +354,10 @@ public abstract class AbstractLauncherUiTest {
 
     // 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 waitForStableState(String message, Supplier<LauncherState> state) {
+    protected void waitForStateTransitionToEnd(String message, Supplier<LauncherState> state) {
         waitForLauncherCondition(message,
-                launcher -> launcher.getStateManager().isInStableState(state.get()));
+                launcher -> launcher.getStateManager().isInStableState(state.get())
+                        && !launcher.getStateManager().isInTransition());
     }
 
     protected void waitForResumed(String message) {
diff --git a/tests/src/com/android/launcher3/ui/WorkProfileTest.java b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
index 2087bfe8c6..aca5951e64 100644
--- a/tests/src/com/android/launcher3/ui/WorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
@@ -91,9 +91,11 @@ public class WorkProfileTest extends AbstractLauncherUiTest {
     public void workTabExists() {
         mDevice.pressHome();
         waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
-        waitForStableState("Launcher internal state didn't switch to Normal", () -> NORMAL);
+        waitForStateTransitionToEnd("Launcher internal state didn't switch to Normal",
+                () -> NORMAL);
         executeOnLauncher(launcher -> launcher.getStateManager().goToState(ALL_APPS));
-        waitForStableState("Launcher internal state didn't switch to All Apps", () -> 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);