diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index c3e46ea..90763b8 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -33,6 +33,7 @@
 
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.tapl.TestHelpers;
+import com.android.systemui.shared.system.QuickStepContract;
 
 import org.junit.Assert;
 import org.junit.rules.TestRule;
@@ -43,6 +44,8 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Test rule that allows executing a test with Quickstep on and then Quickstep off.
@@ -78,11 +81,14 @@
                 @Override
                 public void evaluate() throws Throwable {
                     final Context context = getInstrumentation().getContext();
-                    final String prevOverlayPkg = LauncherInstrumentation.isGesturalMode(context)
-                            ? NAV_BAR_MODE_GESTURAL_OVERLAY
-                            : LauncherInstrumentation.isSwipeUpMode(context)
-                                    ? NAV_BAR_MODE_2BUTTON_OVERLAY
-                                    : NAV_BAR_MODE_3BUTTON_OVERLAY;
+                    final int currentInteractionMode =
+                            LauncherInstrumentation.getCurrentInteractionMode(context);
+                    final String prevOverlayPkg =
+                            QuickStepContract.isGesturalMode(currentInteractionMode)
+                                    ? NAV_BAR_MODE_GESTURAL_OVERLAY
+                                    : QuickStepContract.isSwipeUpMode(currentInteractionMode)
+                                            ? NAV_BAR_MODE_2BUTTON_OVERLAY
+                                            : NAV_BAR_MODE_3BUTTON_OVERLAY;
                     final LauncherInstrumentation.NavigationModel originalMode =
                             mLauncher.getNavigationModel();
                     try {
@@ -131,6 +137,27 @@
                     setOverlayPackageEnabled(NAV_BAR_MODE_GESTURAL_OVERLAY,
                             overlayPackage == NAV_BAR_MODE_GESTURAL_OVERLAY);
 
+                    if (currentSysUiNavigationMode() != expectedMode) {
+                        final CountDownLatch latch = new CountDownLatch(1);
+                        final Context targetContext = getInstrumentation().getTargetContext();
+                        final SysUINavigationMode.NavigationModeChangeListener listener =
+                                newMode -> {
+                                    if (LauncherInstrumentation.getNavigationModel(newMode.resValue)
+                                            == expectedMode) {
+                                        latch.countDown();
+                                    }
+                                };
+                        final SysUINavigationMode sysUINavigationMode =
+                                SysUINavigationMode.INSTANCE.get(targetContext);
+                        targetContext.getMainExecutor().execute(() ->
+                                sysUINavigationMode.addModeChangeListener(listener));
+                        latch.await(10, TimeUnit.SECONDS);
+                        targetContext.getMainExecutor().execute(() ->
+                                sysUINavigationMode.removeModeChangeListener(listener));
+                        Assert.assertTrue("Navigation mode didn't change to " + expectedMode,
+                                currentSysUiNavigationMode() == expectedMode);
+                    }
+
                     for (int i = 0; i != 100; ++i) {
                         if (mLauncher.getNavigationModel() == expectedMode) {
                             Thread.sleep(5000);
@@ -153,4 +180,12 @@
             return base;
         }
     }
+
+    private static LauncherInstrumentation.NavigationModel currentSysUiNavigationMode() {
+        return LauncherInstrumentation.getNavigationModel(
+                SysUINavigationMode.getMode(
+                        getInstrumentation().
+                                getTargetContext()).
+                        resValue);
+    }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index e45fca8..09df396 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -208,14 +208,10 @@
             // app context are not constructed with resources that take overlays into account
             final Context ctx = baseContext.createPackageContext("android", 0);
             for (int i = 0; i < 100; ++i) {
-                log("Interaction mode = " + getCurrentInteractionMode(ctx));
-                if (isGesturalMode(ctx)) {
-                    return NavigationModel.ZERO_BUTTON;
-                } else if (isSwipeUpMode(ctx)) {
-                    return NavigationModel.TWO_BUTTON;
-                } else if (isLegacyMode(ctx)) {
-                    return NavigationModel.THREE_BUTTON;
-                }
+                final int currentInteractionMode = getCurrentInteractionMode(ctx);
+                log("Interaction mode = " + currentInteractionMode);
+                final NavigationModel model = getNavigationModel(currentInteractionMode);
+                if (model != null) return model;
                 Thread.sleep(100);
             }
             fail("Can't detect navigation mode");
@@ -225,6 +221,17 @@
         return NavigationModel.THREE_BUTTON;
     }
 
+    public static NavigationModel getNavigationModel(int currentInteractionMode) {
+        if (QuickStepContract.isGesturalMode(currentInteractionMode)) {
+            return NavigationModel.ZERO_BUTTON;
+        } else if (QuickStepContract.isSwipeUpMode(currentInteractionMode)) {
+            return NavigationModel.TWO_BUTTON;
+        } else if (QuickStepContract.isLegacyMode(currentInteractionMode)) {
+            return NavigationModel.THREE_BUTTON;
+        }
+        return null;
+    }
+
     public static boolean isAvd() {
         return Build.MODEL.contains("Cuttlefish");
     }
@@ -742,19 +749,7 @@
         return currentTime;
     }
 
-    public static boolean isGesturalMode(Context context) {
-        return QuickStepContract.isGesturalMode(getCurrentInteractionMode(context));
-    }
-
-    public static boolean isSwipeUpMode(Context context) {
-        return QuickStepContract.isSwipeUpMode(getCurrentInteractionMode(context));
-    }
-
-    public static boolean isLegacyMode(Context context) {
-        return QuickStepContract.isLegacyMode(getCurrentInteractionMode(context));
-    }
-
-    private static int getCurrentInteractionMode(Context context) {
+    public static int getCurrentInteractionMode(Context context) {
         return getSystemIntegerRes(context, "config_navBarInteractionMode");
     }