diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
index 6cbe171..25adb62 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
@@ -358,14 +358,11 @@
         // Debug if we need to goHome to prevent wrong previous state b/315525621
         mLauncher.goHome();
         assumeFalse(FeatureFlags.ENABLE_BACK_SWIPE_LAUNCHER_ANIMATION.get());
-        mLauncher.getWorkspace().switchToAllApps();
-        mLauncher.pressBack();
-        mLauncher.getWorkspace();
+        mLauncher.getWorkspace().switchToAllApps().pressBackToWorkspace();
         waitForState("Launcher internal state didn't switch to Home", () -> LauncherState.NORMAL);
 
         startAppFast(CALCULATOR_APP_PACKAGE);
-        mLauncher.pressBack();
-        mLauncher.getWorkspace();
+        mLauncher.getLaunchedAppState().pressBackToWorkspace();
         waitForState("Launcher internal state didn't switch to Home", () -> LauncherState.NORMAL);
     }
 
diff --git a/tests/src/com/android/launcher3/allapps/TaplOpenCloseAllApps.java b/tests/src/com/android/launcher3/allapps/TaplOpenCloseAllApps.java
index b4a5169..4f10287 100644
--- a/tests/src/com/android/launcher3/allapps/TaplOpenCloseAllApps.java
+++ b/tests/src/com/android/launcher3/allapps/TaplOpenCloseAllApps.java
@@ -214,8 +214,7 @@
                 .pressBackToWorkspace();
         waitForState("Launcher internal state didn't switch to Home", () -> LauncherState.NORMAL);
         startAppFast(resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR));
-        mLauncher.pressBack();
-        mLauncher.getWorkspace();
+        mLauncher.getLaunchedAppState().pressBackToWorkspace();
         waitForState("Launcher internal state didn't switch to Home", () -> LauncherState.NORMAL);
     }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/LaunchedAppState.java b/tests/tapl/com/android/launcher3/tapl/LaunchedAppState.java
index 184ece7..501c4c3 100644
--- a/tests/tapl/com/android/launcher3/tapl/LaunchedAppState.java
+++ b/tests/tapl/com/android/launcher3/tapl/LaunchedAppState.java
@@ -346,10 +346,14 @@
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
              LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                      "want to press back from launched app to workspace")) {
-            mLauncher.executeAndWaitForWallpaperAnimation(
-                    () -> mLauncher.pressBackImpl(),
-                    "pressing back"
-            );
+            if (mLauncher.isLauncher3()) {
+                mLauncher.pressBackImpl();
+            } else {
+                mLauncher.executeAndWaitForWallpaperAnimation(
+                        () -> mLauncher.pressBackImpl(),
+                        "pressing back"
+                );
+            }
             return new Workspace(mLauncher);
         }
     }