Override DisplayController on main thread.

Bug: 230027385
Flag: TEST_ONLY
Test: TaskbarModeRuleTest
Change-Id: I6e7c349d50c8372ed0d4e1e8d5bfe7f108b2c2e2
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()
             }