diff --git a/quickstep/tests/src/com/android/quickstep/QuickStepOnOffRule.java b/quickstep/tests/src/com/android/quickstep/QuickStepOnOffRule.java
index b801b4f..7274090 100644
--- a/quickstep/tests/src/com/android/quickstep/QuickStepOnOffRule.java
+++ b/quickstep/tests/src/com/android/quickstep/QuickStepOnOffRule.java
@@ -90,11 +90,16 @@
                     base.evaluate();
                 }
 
-                private void overrideSwipeUpEnabled(Boolean swipeUpEnabledOverride) {
+                private void overrideSwipeUpEnabled(Boolean swipeUpEnabledOverride)
+                        throws Throwable {
                     mLauncher.overrideSwipeUpEnabled(swipeUpEnabledOverride);
                     mMainThreadExecutor.execute(() -> OverviewInteractionState.INSTANCE.get(
                             InstrumentationRegistry.getInstrumentation().getTargetContext()).
                             notifySwipeUpSettingChanged(mLauncher.isSwipeUpEnabled()));
+                    // TODO(b/124236673): avoid using sleep().
+                    mLauncher.getDevice().waitForIdle();
+                    Thread.sleep(2000);
+                    mLauncher.getDevice().waitForIdle();
                 }
             };
         } else {
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index e20a89e..b050993 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -402,7 +402,7 @@
     }
 
     @NonNull
-    UiDevice getDevice() {
+    public UiDevice getDevice() {
         return mDevice;
     }