diff --git a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
index 7b73be7ac4..746f8bb98e 100644
--- a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
@@ -22,6 +22,7 @@ import android.content.Intent
 import android.content.pm.PackageManager
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import android.view.Display.DEFAULT_DISPLAY
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
@@ -52,6 +53,7 @@ import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.mockito.Mockito.`when`
 import org.mockito.kotlin.any
@@ -64,8 +66,11 @@ import org.mockito.kotlin.whenever
 import org.mockito.quality.Strictness
 
 /** Test for [DesktopSystemShortcut] */
+// TODO(b/403558856): Improve test coverage for DesktopModeCompatPolicy integration.
 class DesktopSystemShortcutTest {
 
+    @get:Rule val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
+
     private val launcher: RecentsViewContainer = mock()
     private val statsLogManager: StatsLogManager = mock()
     private val statsLogger: StatsLogManager.StatsLogger = mock()