Merge changes I6e7c349d,If6c74b3b,Ib276e9dc,Ida0f7cc0 into main

* changes:
  Override DisplayController on main thread.
  Add annotations for manipulating secure settings.
  Don't use UiThreadTest for Taskbar Unit tests.
  Suspend Launcher taskbar while removed for tests.
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 5a8a1a3..6c3b4ad 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -43,6 +43,8 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_VOICE_INTERACTION_WINDOW_SHOWING;
 import static com.android.wm.shell.Flags.enableTinyTaskbar;
 
+import static java.lang.invoke.MethodHandles.Lookup.PROTECTED;
+
 import android.animation.AnimatorSet;
 import android.animation.ValueAnimator;
 import android.app.ActivityOptions;
@@ -1515,7 +1517,8 @@
         return mIsNavBarKidsMode && isThreeButtonNav();
     }
 
-    protected boolean isNavBarForceVisible() {
+    @VisibleForTesting(otherwise = PROTECTED)
+    public boolean isNavBarForceVisible() {
         return mIsNavBarForceVisible;
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index ee79fbf..b90e5fd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -115,6 +115,7 @@
     private WindowManager mWindowManager;
     private FrameLayout mTaskbarRootLayout;
     private boolean mAddedWindow;
+    private boolean mIsSuspended;
     private final TaskbarNavButtonController mNavButtonController;
     private final ComponentCallbacks mComponentCallbacks;
 
@@ -443,6 +444,8 @@
      */
     @VisibleForTesting
     public synchronized void recreateTaskbar() {
+        if (mIsSuspended) return;
+
         Trace.beginSection("recreateTaskbar");
         try {
             DeviceProfile dp = mUserUnlocked ?
@@ -648,8 +651,22 @@
         }
     }
 
+    /**
+     * Removes Taskbar from the window manager and prevents recreation if {@code true}.
+     * <p>
+     * Suspending is for testing purposes only; avoid calling this method in production.
+     */
     @VisibleForTesting
-    public void addTaskbarRootViewToWindow() {
+    public void setSuspended(boolean isSuspended) {
+        mIsSuspended = isSuspended;
+        if (mIsSuspended) {
+            removeTaskbarRootViewFromWindow();
+        } else {
+            addTaskbarRootViewToWindow();
+        }
+    }
+
+    private void addTaskbarRootViewToWindow() {
         if (enableTaskbarNoRecreate() && !mAddedWindow && mTaskbarActivityContext != null) {
             mWindowManager.addView(mTaskbarRootLayout,
                     mTaskbarActivityContext.getWindowLayoutParams());
@@ -657,8 +674,7 @@
         }
     }
 
-    @VisibleForTesting
-    public void removeTaskbarRootViewFromWindow() {
+    private void removeTaskbarRootViewFromWindow() {
         if (enableTaskbarNoRecreate() && mAddedWindow) {
             mWindowManager.removeViewImmediate(mTaskbarRootLayout);
             mAddedWindow = false;
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
index 9ecd935..2f0b446 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
@@ -20,7 +20,6 @@
 import android.content.ComponentName
 import android.content.Intent
 import android.os.Process
-import androidx.test.annotation.UiThreadTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.BubbleTextView
 import com.android.launcher3.appprediction.PredictionRowView
@@ -34,6 +33,7 @@
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.util.TestUtil
 import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
 import org.junit.Test
@@ -55,17 +55,17 @@
     @InjectController lateinit var overlayController: TaskbarOverlayController
 
     @Test
-    @UiThreadTest
     fun testToggle_once_showsAllApps() {
-        allAppsController.toggle()
+        getInstrumentation().runOnMainSync { allAppsController.toggle() }
         assertThat(allAppsController.isOpen).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testToggle_twice_closesAllApps() {
-        allAppsController.toggle()
-        allAppsController.toggle()
+        getInstrumentation().runOnMainSync {
+            allAppsController.toggle()
+            allAppsController.toggle()
+        }
         assertThat(allAppsController.isOpen).isFalse()
     }
 
@@ -77,54 +77,62 @@
     }
 
     @Test
-    @UiThreadTest
     fun testSetApps_beforeOpened_cachesInfo() {
-        allAppsController.setApps(TEST_APPS, 0, emptyMap())
-        allAppsController.toggle()
+        val overlayContext =
+            TestUtil.getOnUiThread {
+                allAppsController.setApps(TEST_APPS, 0, emptyMap())
+                allAppsController.toggle()
+                overlayController.requestWindow()
+            }
 
-        val overlayContext = overlayController.requestWindow()
         assertThat(overlayContext.appsView.appsStore.apps).isEqualTo(TEST_APPS)
     }
 
     @Test
-    @UiThreadTest
     fun testSetApps_afterOpened_updatesStore() {
-        allAppsController.toggle()
-        allAppsController.setApps(TEST_APPS, 0, emptyMap())
+        val overlayContext =
+            TestUtil.getOnUiThread {
+                allAppsController.toggle()
+                allAppsController.setApps(TEST_APPS, 0, emptyMap())
+                overlayController.requestWindow()
+            }
 
-        val overlayContext = overlayController.requestWindow()
         assertThat(overlayContext.appsView.appsStore.apps).isEqualTo(TEST_APPS)
     }
 
     @Test
-    @UiThreadTest
     fun testSetPredictedApps_beforeOpened_cachesInfo() {
-        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-        allAppsController.toggle()
-
         val predictedApps =
-            overlayController
-                .requestWindow()
-                .appsView
-                .floatingHeaderView
-                .findFixedRowByType(PredictionRowView::class.java)
-                .predictedApps
+            TestUtil.getOnUiThread {
+                allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+                allAppsController.toggle()
+
+                overlayController
+                    .requestWindow()
+                    .appsView
+                    .floatingHeaderView
+                    .findFixedRowByType(PredictionRowView::class.java)
+                    .predictedApps
+            }
+
         assertThat(predictedApps).isEqualTo(TEST_PREDICTED_APPS)
     }
 
     @Test
-    @UiThreadTest
     fun testSetPredictedApps_afterOpened_cachesInfo() {
-        allAppsController.toggle()
-        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-
         val predictedApps =
-            overlayController
-                .requestWindow()
-                .appsView
-                .floatingHeaderView
-                .findFixedRowByType(PredictionRowView::class.java)
-                .predictedApps
+            TestUtil.getOnUiThread {
+                allAppsController.toggle()
+                allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+
+                overlayController
+                    .requestWindow()
+                    .appsView
+                    .floatingHeaderView
+                    .findFixedRowByType(PredictionRowView::class.java)
+                    .predictedApps
+            }
+
         assertThat(predictedApps).isEqualTo(TEST_PREDICTED_APPS)
     }
 
@@ -140,36 +148,38 @@
         }
 
         // Ensure the recycler view fully inflates before trying to grab an icon.
-        getInstrumentation().runOnMainSync {
-            val btv =
+        val btv =
+            TestUtil.getOnUiThread {
                 overlayController
                     .requestWindow()
                     .appsView
                     .activeRecyclerView
                     .findViewHolderForAdapterPosition(0)
                     ?.itemView as? BubbleTextView
-            assertThat(btv?.hasDot()).isTrue()
-        }
+            }
+        assertThat(btv?.hasDot()).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testUpdateNotificationDots_predictedApp_hasDot() {
-        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-        allAppsController.toggle()
+        getInstrumentation().runOnMainSync {
+            allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+            allAppsController.toggle()
+            taskbarUnitTestRule.activityContext.popupDataProvider.onNotificationPosted(
+                PackageUserKey.fromItemInfo(TEST_PREDICTED_APPS[0]),
+                NotificationKeyData("key"),
+            )
+        }
 
-        taskbarUnitTestRule.activityContext.popupDataProvider.onNotificationPosted(
-            PackageUserKey.fromItemInfo(TEST_PREDICTED_APPS[0]),
-            NotificationKeyData("key"),
-        )
-
-        val predictionRowView =
-            overlayController
-                .requestWindow()
-                .appsView
-                .floatingHeaderView
-                .findFixedRowByType(PredictionRowView::class.java)
-        val btv = predictionRowView.getChildAt(0) as BubbleTextView
+        val btv =
+            TestUtil.getOnUiThread {
+                overlayController
+                    .requestWindow()
+                    .appsView
+                    .floatingHeaderView
+                    .findFixedRowByType(PredictionRowView::class.java)
+                    .getChildAt(0) as BubbleTextView
+            }
         assertThat(btv.hasDot()).isTrue()
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
index fae5562..f946d4d 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
@@ -18,7 +18,6 @@
 
 import android.app.ActivityManager.RunningTaskInfo
 import android.view.MotionEvent
-import androidx.test.annotation.UiThreadTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.AbstractFloatingView
 import com.android.launcher3.AbstractFloatingView.TYPE_OPTIONS_POPUP
@@ -31,7 +30,7 @@
 import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
-import com.android.launcher3.views.BaseDragLayer
+import com.android.launcher3.util.TestUtil.getOnUiThread
 import com.android.systemui.shared.system.TaskStackChangeListeners
 import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
@@ -54,74 +53,69 @@
         get() = taskbarUnitTestRule.activityContext
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_twice_reusesWindow() {
-        val context1 = overlayController.requestWindow()
-        val context2 = overlayController.requestWindow()
+        val (context1, context2) =
+            getOnUiThread {
+                Pair(overlayController.requestWindow(), overlayController.requestWindow())
+            }
         assertThat(context1).isSameInstanceAs(context2)
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_afterHidingExistingWindow_createsNewWindow() {
-        val context1 = overlayController.requestWindow()
-        overlayController.hideWindow()
+        val context1 = getOnUiThread { overlayController.requestWindow() }
+        getInstrumentation().runOnMainSync { overlayController.hideWindow() }
 
-        val context2 = overlayController.requestWindow()
+        val context2 = getOnUiThread { overlayController.requestWindow() }
         assertThat(context1).isNotSameInstanceAs(context2)
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_afterHidingOverlay_createsNewWindow() {
-        val context1 = overlayController.requestWindow()
-        TestOverlayView.show(context1)
-        overlayController.hideWindow()
+        val context1 = getOnUiThread { overlayController.requestWindow() }
+        getInstrumentation().runOnMainSync {
+            TestOverlayView.show(context1)
+            overlayController.hideWindow()
+        }
 
-        val context2 = overlayController.requestWindow()
+        val context2 = getOnUiThread { overlayController.requestWindow() }
         assertThat(context1).isNotSameInstanceAs(context2)
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_addsProxyView() {
-        TestOverlayView.show(overlayController.requestWindow())
+        getInstrumentation().runOnMainSync {
+            TestOverlayView.show(overlayController.requestWindow())
+        }
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_closeProxyView_closesOverlay() {
-        val overlay = TestOverlayView.show(overlayController.requestWindow())
-        AbstractFloatingView.closeOpenContainer(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
+        getInstrumentation().runOnMainSync {
+            AbstractFloatingView.closeOpenContainer(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)
+        }
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
     fun testRequestWindow_attachesDragLayer() {
-        lateinit var dragLayer: BaseDragLayer<*>
-        getInstrumentation().runOnMainSync {
-            dragLayer = overlayController.requestWindow().dragLayer
-        }
-
+        val dragLayer = getOnUiThread { overlayController.requestWindow().dragLayer }
         // Allow drag layer to attach before checking.
         getInstrumentation().runOnMainSync { assertThat(dragLayer.isAttachedToWindow).isTrue() }
     }
 
     @Test
-    @UiThreadTest
     fun testHideWindow_closesOverlay() {
-        val overlay = TestOverlayView.show(overlayController.requestWindow())
-        overlayController.hideWindow()
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
+        getInstrumentation().runOnMainSync { overlayController.hideWindow() }
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
     fun testHideWindow_detachesDragLayer() {
-        lateinit var dragLayer: BaseDragLayer<*>
-        getInstrumentation().runOnMainSync {
-            dragLayer = overlayController.requestWindow().dragLayer
-        }
+        val dragLayer = getOnUiThread { overlayController.requestWindow().dragLayer }
 
         // Wait for drag layer to be attached to window before hiding.
         getInstrumentation().runOnMainSync {
@@ -131,26 +125,30 @@
     }
 
     @Test
-    @UiThreadTest
     fun testTwoOverlays_closeOne_windowStaysOpen() {
-        val context = overlayController.requestWindow()
-        val overlay1 = TestOverlayView.show(context)
-        val overlay2 = TestOverlayView.show(context)
+        val (overlay1, overlay2) =
+            getOnUiThread {
+                val context = overlayController.requestWindow()
+                Pair(TestOverlayView.show(context), TestOverlayView.show(context))
+            }
 
-        overlay1.close(false)
+        getInstrumentation().runOnMainSync { overlay1.close(false) }
         assertThat(overlay2.isOpen).isTrue()
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testTwoOverlays_closeAll_closesWindow() {
-        val context = overlayController.requestWindow()
-        val overlay1 = TestOverlayView.show(context)
-        val overlay2 = TestOverlayView.show(context)
+        val (overlay1, overlay2) =
+            getOnUiThread {
+                val context = overlayController.requestWindow()
+                Pair(TestOverlayView.show(context), TestOverlayView.show(context))
+            }
 
-        overlay1.close(false)
-        overlay2.close(false)
+        getInstrumentation().runOnMainSync {
+            overlay1.close(false)
+            overlay2.close(false)
+        }
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isFalse()
     }
 
@@ -165,11 +163,7 @@
 
     @Test
     fun testTaskMovedToFront_closesOverlay() {
-        lateinit var overlay: TestOverlayView
-        getInstrumentation().runOnMainSync {
-            overlay = TestOverlayView.show(overlayController.requestWindow())
-        }
-
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
         TaskStackChangeListeners.getInstance().listenerImpl.onTaskMovedToFront(RunningTaskInfo())
         // Make sure TaskStackChangeListeners' Handler posts the callback before checking state.
         getInstrumentation().runOnMainSync { assertThat(overlay.isOpen).isFalse() }
@@ -177,9 +171,8 @@
 
     @Test
     fun testTaskStackChanged_allAppsClosed_overlayStaysOpen() {
-        lateinit var overlay: TestOverlayView
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
         getInstrumentation().runOnMainSync {
-            overlay = TestOverlayView.show(overlayController.requestWindow())
             taskbarContext.controllers.sharedState?.allAppsVisible = false
         }
 
@@ -189,9 +182,8 @@
 
     @Test
     fun testTaskStackChanged_allAppsOpen_closesOverlay() {
-        lateinit var overlay: TestOverlayView
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
         getInstrumentation().runOnMainSync {
-            overlay = TestOverlayView.show(overlayController.requestWindow())
             taskbarContext.controllers.sharedState?.allAppsVisible = true
         }
 
@@ -200,33 +192,39 @@
     }
 
     @Test
-    @UiThreadTest
     fun testUpdateLauncherDeviceProfile_overlayNotRebindSafe_closesOverlay() {
-        val overlayContext = overlayController.requestWindow()
-        val overlay = TestOverlayView.show(overlayContext).apply { type = TYPE_OPTIONS_POPUP }
+        val context = getOnUiThread { overlayController.requestWindow() }
+        val overlay = getOnUiThread {
+            TestOverlayView.show(context).apply { type = TYPE_OPTIONS_POPUP }
+        }
 
-        overlayController.updateLauncherDeviceProfile(
-            overlayController.launcherDeviceProfile
-                .toBuilder(overlayContext)
-                .setGestureMode(false)
-                .build()
-        )
+        getInstrumentation().runOnMainSync {
+            overlayController.updateLauncherDeviceProfile(
+                overlayController.launcherDeviceProfile
+                    .toBuilder(context)
+                    .setGestureMode(false)
+                    .build()
+            )
+        }
 
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
-    @UiThreadTest
     fun testUpdateLauncherDeviceProfile_overlayRebindSafe_overlayStaysOpen() {
-        val overlayContext = overlayController.requestWindow()
-        val overlay = TestOverlayView.show(overlayContext).apply { type = TYPE_TASKBAR_ALL_APPS }
+        val context = getOnUiThread { overlayController.requestWindow() }
+        val overlay = getOnUiThread {
+            TestOverlayView.show(context).apply { type = TYPE_TASKBAR_ALL_APPS }
+        }
 
-        overlayController.updateLauncherDeviceProfile(
-            overlayController.launcherDeviceProfile
-                .toBuilder(overlayContext)
-                .setGestureMode(false)
-                .build()
-        )
+        getInstrumentation().runOnMainSync {
+            overlayController.updateLauncherDeviceProfile(
+                overlayController.launcherDeviceProfile
+                    .toBuilder(context)
+                    .setGestureMode(false)
+                    .build()
+            )
+        }
 
         assertThat(overlay.isOpen).isTrue()
     }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
index 6638736..c48947e 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.taskbar.rules
 
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
 import com.android.launcher3.util.DisplayController
@@ -59,23 +60,25 @@
             override fun evaluate() {
                 val mode = taskbarMode.mode
 
-                context.applicationContext.putObject(
-                    DisplayController.INSTANCE,
-                    object : DisplayController(context) {
-                        override fun getInfo(): Info {
-                            return spy(super.getInfo()) {
-                                on { isTransientTaskbar } doReturn (mode == Mode.TRANSIENT)
-                                on { isPinnedTaskbar } doReturn (mode == Mode.PINNED)
-                                on { navigationMode } doReturn
-                                    when (mode) {
-                                        Mode.TRANSIENT,
-                                        Mode.PINNED -> NavigationMode.NO_BUTTON
-                                        Mode.THREE_BUTTONS -> NavigationMode.THREE_BUTTONS
-                                    }
+                getInstrumentation().runOnMainSync {
+                    context.applicationContext.putObject(
+                        DisplayController.INSTANCE,
+                        object : DisplayController(context) {
+                            override fun getInfo(): Info {
+                                return spy(super.getInfo()) {
+                                    on { isTransientTaskbar } doReturn (mode == Mode.TRANSIENT)
+                                    on { isPinnedTaskbar } doReturn (mode == Mode.PINNED)
+                                    on { navigationMode } doReturn
+                                        when (mode) {
+                                            Mode.TRANSIENT,
+                                            Mode.PINNED -> NavigationMode.NO_BUTTON
+                                            Mode.THREE_BUTTONS -> NavigationMode.THREE_BUTTONS
+                                        }
+                                }
                             }
-                        }
-                    },
-                )
+                        },
+                    )
+                }
 
                 base.evaluate()
             }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
index 12f946e..8a64949 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
@@ -20,18 +20,25 @@
 import android.app.PendingIntent
 import android.content.IIntentSender
 import android.content.Intent
+import android.provider.Settings
+import android.provider.Settings.Secure.NAV_BAR_KIDS_MODE
+import android.provider.Settings.Secure.USER_SETUP_COMPLETE
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.rule.ServiceTestRule
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.taskbar.TaskbarActivityContext
 import com.android.launcher3.taskbar.TaskbarManager
 import com.android.launcher3.taskbar.TaskbarNavButtonController.TaskbarNavButtonCallbacks
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
 import com.android.launcher3.util.LauncherMultivalentJUnit.Companion.isRunningInRobolectric
+import com.android.launcher3.util.ModelTestExtensions.loadModelSync
+import com.android.launcher3.util.TestUtil
 import com.android.quickstep.AllAppsActionManager
 import com.android.quickstep.TouchInteractionService
 import com.android.quickstep.TouchInteractionService.TISBinder
 import org.junit.Assume.assumeTrue
+import org.junit.rules.RuleChain
 import org.junit.rules.TestRule
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
@@ -48,12 +55,11 @@
  * that code that is executed on the main thread in production should also happen on that thread
  * when tested.
  *
- * `@UiThreadTest` is a simple way to run an entire test body on the main thread. But if a test
- * executes code that appends message(s) to the main thread's `MessageQueue`, the annotation will
- * prevent those messages from being processed until after the test body finishes.
+ * `@UiThreadTest` is incompatible with this rule. The annotation causes this rule to run on the
+ * main thread, but it needs to be run on the test thread for it to work properly. Instead, only run
+ * code that requires the main thread using something like [Instrumentation.runOnMainSync] or
+ * [TestUtil.getOnUiThread].
  *
- * To test pending messages, instead use something like [Instrumentation.runOnMainSync] to perform
- * only sections of the test body on the main thread synchronously:
  * ```
  * @Test
  * fun example() {
@@ -71,6 +77,10 @@
     private val instrumentation = InstrumentationRegistry.getInstrumentation()
     private val serviceTestRule = ServiceTestRule()
 
+    private val userSetupCompleteRule = TaskbarSecureSettingRule(USER_SETUP_COMPLETE)
+    private val kidsModeRule = TaskbarSecureSettingRule(NAV_BAR_KIDS_MODE)
+    private val settingRules = RuleChain.outerRule(userSetupCompleteRule).around(kidsModeRule)
+
     private lateinit var taskbarManager: TaskbarManager
 
     val activityContext: TaskbarActivityContext
@@ -80,15 +90,34 @@
         }
 
     override fun apply(base: Statement, description: Description): Statement {
+        return settingRules.apply(createStatement(base, description), description)
+    }
+
+    private fun createStatement(base: Statement, description: Description): Statement {
         return object : Statement() {
             override fun evaluate() {
 
+                // Only run test when Taskbar is enabled.
                 instrumentation.runOnMainSync {
                     assumeTrue(
                         LauncherAppState.getIDP(context).getDeviceProfile(context).isTaskbarPresent
                     )
                 }
 
+                // Process secure setting annotations.
+                instrumentation.runOnMainSync {
+                    userSetupCompleteRule.putInt(
+                        if (description.getAnnotation(UserSetupMode::class.java) != null) {
+                            0
+                        } else {
+                            1
+                        }
+                    )
+                    kidsModeRule.putInt(
+                        if (description.getAnnotation(NavBarKidsMode::class.java) != null) 1 else 0
+                    )
+                }
+
                 // Check for existing Taskbar instance from Launcher process.
                 val launcherTaskbarManager: TaskbarManager? =
                     if (!isRunningInRobolectric) {
@@ -105,8 +134,8 @@
                         null
                     }
 
-                instrumentation.runOnMainSync {
-                    taskbarManager =
+                taskbarManager =
+                    TestUtil.getOnUiThread {
                         TaskbarManager(
                             context,
                             AllAppsActionManager(context, UI_HELPER_EXECUTOR) {
@@ -114,12 +143,14 @@
                             },
                             object : TaskbarNavButtonCallbacks {},
                         )
-                }
+                    }
 
                 try {
+                    LauncherAppState.getInstance(context).model.loadModelSync()
+
                     // Replace Launcher Taskbar window with test instance.
                     instrumentation.runOnMainSync {
-                        launcherTaskbarManager?.removeTaskbarRootViewFromWindow()
+                        launcherTaskbarManager?.setSuspended(true)
                         taskbarManager.onUserUnlocked() // Required to complete initialization.
                     }
 
@@ -129,7 +160,7 @@
                     // Revert Taskbar window.
                     instrumentation.runOnMainSync {
                         taskbarManager.destroy()
-                        launcherTaskbarManager?.addTaskbarRootViewToWindow()
+                        launcherTaskbarManager?.setSuspended(false)
                     }
                 }
             }
@@ -167,4 +198,35 @@
     @Retention(AnnotationRetention.RUNTIME)
     @Target(AnnotationTarget.FIELD)
     annotation class InjectController
+
+    /** Overrides [USER_SETUP_COMPLETE] to be `false` for tests. */
+    @Retention(AnnotationRetention.RUNTIME)
+    @Target(AnnotationTarget.CLASS, AnnotationTarget.FUNCTION)
+    annotation class UserSetupMode
+
+    /** Overrides [NAV_BAR_KIDS_MODE] to be `true` for tests. */
+    @Retention(AnnotationRetention.RUNTIME)
+    @Target(AnnotationTarget.CLASS, AnnotationTarget.FUNCTION)
+    annotation class NavBarKidsMode
+
+    /** Rule for Taskbar integer-based secure settings. */
+    private inner class TaskbarSecureSettingRule(private val settingName: String) : TestRule {
+
+        override fun apply(base: Statement, description: Description): Statement {
+            return object : Statement() {
+                override fun evaluate() {
+                    val originalValue =
+                        Settings.Secure.getInt(context.contentResolver, settingName, /* def= */ 0)
+                    try {
+                        base.evaluate()
+                    } finally {
+                        instrumentation.runOnMainSync { putInt(originalValue) }
+                    }
+                }
+            }
+        }
+
+        /** Puts [value] into secure settings under [settingName]. */
+        fun putInt(value: Int) = Settings.Secure.putInt(context.contentResolver, settingName, value)
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
index 8262e0f..234e499 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
@@ -22,6 +22,8 @@
 import com.android.launcher3.taskbar.TaskbarManager
 import com.android.launcher3.taskbar.TaskbarStashController
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.NavBarKidsMode
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.UserSetupMode
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.google.common.truth.Truth.assertThat
@@ -125,9 +127,40 @@
         }
     }
 
-    /** Executes [runTest] after the [testRule] setup phase completes. */
+    @Test
+    fun testUserSetupMode_default_isComplete() {
+        onSetup { assertThat(activityContext.isUserSetupComplete).isTrue() }
+    }
+
+    @Test
+    fun testUserSetupMode_withAnnotation_isIncomplete() {
+        @UserSetupMode class Mode
+        onSetup(description = Description.createSuiteDescription(Mode::class.java)) {
+            assertThat(activityContext.isUserSetupComplete).isFalse()
+        }
+    }
+
+    @Test
+    fun testNavBarKidsMode_default_navBarNotForcedVisible() {
+        onSetup { assertThat(activityContext.isNavBarForceVisible).isFalse() }
+    }
+
+    @Test
+    fun testNavBarKidsMode_withAnnotation_navBarForcedVisible() {
+        @NavBarKidsMode class Mode
+        onSetup(description = Description.createSuiteDescription(Mode::class.java)) {
+            assertThat(activityContext.isNavBarForceVisible).isTrue()
+        }
+    }
+
+    /**
+     * Executes [runTest] after the [testRule] setup phase completes.
+     *
+     * A [description] can also be provided to mimic annotating a test or test class.
+     */
     private fun onSetup(
         testRule: TaskbarUnitTestRule = TaskbarUnitTestRule(this, context),
+        description: Description = DESCRIPTION,
         runTest: TaskbarUnitTestRule.() -> Unit,
     ) {
         testRule
@@ -135,7 +168,7 @@
                 object : Statement() {
                     override fun evaluate() = runTest(testRule)
                 },
-                DESCRIPTION,
+                description,
             )
             .evaluate()
     }