diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRule.kt
index 60c8dec..2d26bbf 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRule.kt
@@ -16,11 +16,15 @@
 
 package com.android.launcher3.taskbar.rules
 
+import android.platform.test.flag.junit.FlagsParameterization
+import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.ConstantItem
+import com.android.launcher3.Flags.FLAG_ENABLE_TASKBAR_PINNING
 import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.LauncherPrefs.Companion.TASKBAR_PINNING
 import com.android.launcher3.LauncherPrefs.Companion.TASKBAR_PINNING_IN_DESKTOP_MODE
+import com.android.launcher3.util.DisplayController
 import kotlin.reflect.KProperty
 import org.junit.rules.TestRule
 import org.junit.runner.Description
@@ -40,13 +44,20 @@
 class TaskbarPinningPreferenceRule(context: TaskbarWindowSandboxContext) : TestRule {
 
     private val prefs = LauncherPrefs.get(context)
+    private val setFlagsRule =
+        SetFlagsRule(FlagsParameterization(mapOf(FLAG_ENABLE_TASKBAR_PINNING to true)))
 
     var isPinned by PinningPreference(TASKBAR_PINNING)
     var isPinnedInDesktopMode by PinningPreference(TASKBAR_PINNING_IN_DESKTOP_MODE)
 
     override fun apply(base: Statement, description: Description): Statement {
+        return setFlagsRule.apply(createStatement(base), description)
+    }
+
+    private fun createStatement(base: Statement): Statement {
         return object : Statement() {
             override fun evaluate() {
+                DisplayController.enableTaskbarModePreferenceForTests(true)
                 val wasPinned = isPinned
                 val wasPinnedInDesktopMode = isPinnedInDesktopMode
                 try {
@@ -54,6 +65,7 @@
                 } finally {
                     isPinned = wasPinned
                     isPinnedInDesktopMode = wasPinnedInDesktopMode
+                    DisplayController.enableTaskbarModePreferenceForTests(false)
                 }
             }
         }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt
index 3f0a238..a709133 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt
@@ -19,6 +19,8 @@
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.LauncherMultivalentJUnit
+import com.android.launcher3.util.window.WindowManagerProxy
+import com.google.android.apps.nexuslauncher.deviceemulator.TestWindowManagerProxy
 import com.google.common.truth.Truth.assertThat
 import org.junit.Test
 import org.junit.runner.Description
@@ -51,6 +53,11 @@
 
     @Test
     fun testEnableDesktopPinning_verifyDisplayController() {
+        context.applicationContext.putObject(
+            WindowManagerProxy.INSTANCE,
+            TestWindowManagerProxy(context).apply { isInDesktopMode = true },
+        )
+
         onSetup {
             preferenceRule.isPinned = false
             preferenceRule.isPinnedInDesktopMode = true
@@ -60,6 +67,11 @@
 
     @Test
     fun testDisableDesktopPinning_verifyDisplayController() {
+        context.applicationContext.putObject(
+            WindowManagerProxy.INSTANCE,
+            TestWindowManagerProxy(context).apply { isInDesktopMode = true },
+        )
+
         onSetup {
             preferenceRule.isPinned = false
             preferenceRule.isPinnedInDesktopMode = false
@@ -83,12 +95,14 @@
 
     /** Executes [runTest] after the [preferenceRule] setup phase completes. */
     private fun onSetup(runTest: () -> Unit) {
-        preferenceRule.apply(
-            object : Statement() {
-                override fun evaluate() = runTest()
-            },
-            DESCRIPTION,
-        )
+        preferenceRule
+            .apply(
+                object : Statement() {
+                    override fun evaluate() = runTest()
+                },
+                DESCRIPTION,
+            )
+            .evaluate()
     }
 
     private companion object {
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
index 321e7a9..ee21df8 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
@@ -50,7 +50,7 @@
         )
     }
 
-    override fun getApplicationContext() = application
+    override fun getApplicationContext(): SandboxContext = application
 
     companion object {
         /** Creates a [TaskbarWindowSandboxContext] to sandbox [base] for Taskbar tests. */
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index 3dcc663..5b2f737 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -76,6 +76,7 @@
 
     private static final String TAG = "DisplayController";
     private static final boolean DEBUG = false;
+    private static boolean sTaskbarModePreferenceStatusForTests = false;
     private static boolean sTransientTaskbarStatusForTests = true;
 
     // TODO(b/254119092) remove all logs with this tag
@@ -203,6 +204,14 @@
     }
 
     /**
+     * Enables respecting taskbar mode preference during test.
+     */
+    @VisibleForTesting
+    public static void enableTaskbarModePreferenceForTests(boolean enable) {
+        sTaskbarModePreferenceStatusForTests = enable;
+    }
+
+    /**
      * Returns whether the taskbar is pinned in gesture navigation mode.
      */
     public static boolean isPinnedTaskbar(Context context) {
@@ -461,7 +470,7 @@
             if (navigationMode != NavigationMode.NO_BUTTON) {
                 return false;
             }
-            if (Utilities.isRunningInTestHarness()) {
+            if (Utilities.isRunningInTestHarness() && !sTaskbarModePreferenceStatusForTests) {
                 // TODO(b/258604917): Once ENABLE_TASKBAR_PINNING is enabled, remove usage of
                 //  sTransientTaskbarStatusForTests and update test to directly
                 //  toggle shared preference to switch transient taskbar on/off.