Merge "Move logic to move the shade to the default display once locked to ShadeDisplaysRepository" into main

10 files changed, 135 insertions, 142 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepositoryTest.kt
index 007a0fb87953..4f332d4bbed8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepositoryTest.kt
@@ -17,15 +17,21 @@
 package com.android.systemui.shade.data.repository
 
 import android.provider.Settings.Global.DEVELOPMENT_SHADE_DISPLAY_AWARENESS
+import android.view.Display
+import android.view.Display.TYPE_EXTERNAL
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
+import com.android.systemui.display.data.repository.display
 import com.android.systemui.display.data.repository.displayRepository
+import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.useUnconfinedTestDispatcher
 import com.android.systemui.shade.display.AnyExternalShadeDisplayPolicy
 import com.android.systemui.shade.display.DefaultDisplayShadePolicy
+import com.android.systemui.shade.display.FakeShadeDisplayPolicy
 import com.android.systemui.shade.display.StatusBarTouchShadeDisplayPolicy
 import com.android.systemui.testKosmos
 import com.android.systemui.util.settings.fakeGlobalSettings
@@ -37,24 +43,28 @@ import org.junit.runner.RunWith
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class ShadeDisplaysRepositoryTest : SysuiTestCase() {
-    private val kosmos = testKosmos()
+    private val kosmos = testKosmos().useUnconfinedTestDispatcher()
     private val testScope = kosmos.testScope
     private val globalSettings = kosmos.fakeGlobalSettings
     private val displayRepository = kosmos.displayRepository
     private val defaultPolicy = DefaultDisplayShadePolicy()
     private val policies = kosmos.shadeDisplayPolicies
+    private val keyguardRepository = kosmos.fakeKeyguardRepository
 
-    private val underTest =
+    private fun createUnderTest(shadeOnDefaultDisplayWhenLocked: Boolean = false) =
         ShadeDisplaysRepositoryImpl(
             globalSettings,
             defaultPolicy,
             testScope.backgroundScope,
             policies,
+            shadeOnDefaultDisplayWhenLocked = shadeOnDefaultDisplayWhenLocked,
+            keyguardRepository,
         )
 
     @Test
     fun policy_changing_propagatedFromTheLatestPolicy() =
         testScope.runTest {
+            val underTest = createUnderTest()
             val displayIds by collectValues(underTest.displayId)
 
             assertThat(displayIds).containsExactly(0)
@@ -81,30 +91,54 @@ class ShadeDisplaysRepositoryTest : SysuiTestCase() {
     @Test
     fun policy_updatesBasedOnSettingValue_defaultDisplay() =
         testScope.runTest {
-            val policy by collectLastValue(underTest.policy)
-
+            val underTest = createUnderTest()
             globalSettings.putString(DEVELOPMENT_SHADE_DISPLAY_AWARENESS, "default_display")
 
-            assertThat(policy).isInstanceOf(DefaultDisplayShadePolicy::class.java)
+            assertThat(underTest.currentPolicy).isInstanceOf(DefaultDisplayShadePolicy::class.java)
         }
 
     @Test
     fun policy_updatesBasedOnSettingValue_anyExternal() =
         testScope.runTest {
-            val policy by collectLastValue(underTest.policy)
-
+            val underTest = createUnderTest()
             globalSettings.putString(DEVELOPMENT_SHADE_DISPLAY_AWARENESS, "any_external_display")
 
-            assertThat(policy).isInstanceOf(AnyExternalShadeDisplayPolicy::class.java)
+            assertThat(underTest.currentPolicy)
+                .isInstanceOf(AnyExternalShadeDisplayPolicy::class.java)
         }
 
     @Test
     fun policy_updatesBasedOnSettingValue_focusBased() =
         testScope.runTest {
-            val policy by collectLastValue(underTest.policy)
-
+            val underTest = createUnderTest()
             globalSettings.putString(DEVELOPMENT_SHADE_DISPLAY_AWARENESS, "status_bar_latest_touch")
 
-            assertThat(policy).isInstanceOf(StatusBarTouchShadeDisplayPolicy::class.java)
+            assertThat(underTest.currentPolicy)
+                .isInstanceOf(StatusBarTouchShadeDisplayPolicy::class.java)
+        }
+
+    @Test
+    fun displayId_afterKeyguardHides_goesBackToPreviousDisplay() =
+        testScope.runTest {
+            val underTest = createUnderTest(shadeOnDefaultDisplayWhenLocked = true)
+            globalSettings.putString(
+                DEVELOPMENT_SHADE_DISPLAY_AWARENESS,
+                FakeShadeDisplayPolicy.name,
+            )
+
+            val displayId by collectLastValue(underTest.displayId)
+
+            displayRepository.addDisplays(display(id = 2, type = TYPE_EXTERNAL))
+            FakeShadeDisplayPolicy.setDisplayId(2)
+
+            assertThat(displayId).isEqualTo(2)
+
+            keyguardRepository.setKeyguardShowing(true)
+
+            assertThat(displayId).isEqualTo(Display.DEFAULT_DISPLAY)
+
+            keyguardRepository.setKeyguardShowing(false)
+
+            assertThat(displayId).isEqualTo(2)
         }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/data/repository/ShadePrimaryDisplayCommandTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/data/repository/ShadePrimaryDisplayCommandTest.kt
index eeb3e6b31c69..fd6bc98b006c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/data/repository/ShadePrimaryDisplayCommandTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/data/repository/ShadePrimaryDisplayCommandTest.kt
@@ -25,7 +25,7 @@ import com.android.systemui.display.data.repository.displayRepository
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.kosmos.useUnconfinedTestDispatcher
 import com.android.systemui.shade.ShadePrimaryDisplayCommand
-import com.android.systemui.shade.display.ShadeDisplayPolicy
+import com.android.systemui.shade.display.FakeShadeDisplayPolicy
 import com.android.systemui.statusbar.commandline.commandRegistry
 import com.android.systemui.testKosmos
 import com.android.systemui.util.settings.fakeGlobalSettings
@@ -33,8 +33,6 @@ import com.google.common.truth.StringSubject
 import com.google.common.truth.Truth.assertThat
 import java.io.PrintWriter
 import java.io.StringWriter
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
@@ -118,23 +116,12 @@ class ShadePrimaryDisplayCommandTest : SysuiTestCase() {
     @Test
     fun policies_setsNewPolicy() =
         testScope.runTest {
-            val policy by collectLastValue(shadeDisplaysRepository.policy)
-            val newPolicy = policies.last().name
+            val newPolicy = FakeShadeDisplayPolicy.name
 
             commandRegistry.onShellCommand(pw, arrayOf("shade_display_override", newPolicy))
 
-            assertThat(policy!!.name).isEqualTo(newPolicy)
+            assertThat(shadeDisplaysRepository.currentPolicy.name).isEqualTo(newPolicy)
         }
-
-    private fun makePolicy(policyName: String): ShadeDisplayPolicy {
-        return object : ShadeDisplayPolicy {
-            override val name: String
-                get() = policyName
-
-            override val displayId: StateFlow<Int>
-                get() = MutableStateFlow(0)
-        }
-    }
 }
 
 private fun StringSubject.containsAllIn(strings: List<String>) {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicyTest.kt
index 20dfd3e11947..e43c46b36a06 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicyTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicyTest.kt
@@ -26,12 +26,11 @@ import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
 import com.android.systemui.display.data.repository.display
 import com.android.systemui.display.data.repository.displayRepository
-import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.kosmos.useUnconfinedTestDispatcher
+import com.android.systemui.shade.data.repository.statusBarTouchShadeDisplayPolicy
 import com.android.systemui.shade.domain.interactor.notificationElement
 import com.android.systemui.shade.domain.interactor.qsElement
-import com.android.systemui.shade.domain.interactor.shadeInteractor
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlin.test.Test
@@ -45,22 +44,9 @@ import org.mockito.kotlin.mock
 class StatusBarTouchShadeDisplayPolicyTest : SysuiTestCase() {
     private val kosmos = testKosmos().useUnconfinedTestDispatcher()
     private val testScope = kosmos.testScope
-    private val keyguardRepository = kosmos.fakeKeyguardRepository
     private val displayRepository = kosmos.displayRepository
 
-    private fun createUnderTest(
-        shadeOnDefaultDisplayWhenLocked: Boolean = false
-    ): StatusBarTouchShadeDisplayPolicy {
-        return StatusBarTouchShadeDisplayPolicy(
-            displayRepository,
-            keyguardRepository,
-            testScope.backgroundScope,
-            shadeOnDefaultDisplayWhenLocked = shadeOnDefaultDisplayWhenLocked,
-            shadeInteractor = { kosmos.shadeInteractor },
-            { kosmos.qsElement },
-            { kosmos.notificationElement },
-        )
-    }
+    private val underTest = kosmos.statusBarTouchShadeDisplayPolicy
 
     private fun createMotionEventForDisplay(displayId: Int, xCoordinate: Float = 0f): MotionEvent {
         return mock<MotionEvent> {
@@ -71,15 +57,12 @@ class StatusBarTouchShadeDisplayPolicyTest : SysuiTestCase() {
 
     @Test
     fun displayId_defaultToDefaultDisplay() {
-        val underTest = createUnderTest()
-
         assertThat(underTest.displayId.value).isEqualTo(Display.DEFAULT_DISPLAY)
     }
 
     @Test
     fun onStatusBarTouched_called_updatesDisplayId() =
         testScope.runTest {
-            val underTest = createUnderTest()
             val displayId by collectLastValue(underTest.displayId)
 
             displayRepository.addDisplays(display(id = 2, type = TYPE_EXTERNAL))
@@ -91,7 +74,6 @@ class StatusBarTouchShadeDisplayPolicyTest : SysuiTestCase() {
     @Test
     fun onStatusBarTouched_notExistentDisplay_displayIdNotUpdated() =
         testScope.runTest {
-            val underTest = createUnderTest()
             val displayIds by collectValues(underTest.displayId)
             assertThat(displayIds).isEqualTo(listOf(Display.DEFAULT_DISPLAY))
 
@@ -104,7 +86,6 @@ class StatusBarTouchShadeDisplayPolicyTest : SysuiTestCase() {
     @Test
     fun onStatusBarTouched_afterDisplayRemoved_goesBackToDefaultDisplay() =
         testScope.runTest {
-            val underTest = createUnderTest()
             val displayId by collectLastValue(underTest.displayId)
 
             displayRepository.addDisplays(display(id = 2, type = TYPE_EXTERNAL))
@@ -118,46 +99,8 @@ class StatusBarTouchShadeDisplayPolicyTest : SysuiTestCase() {
         }
 
     @Test
-    fun onStatusBarTouched_afterKeyguardVisible_goesBackToDefaultDisplay() =
-        testScope.runTest {
-            val underTest = createUnderTest(shadeOnDefaultDisplayWhenLocked = true)
-            val displayId by collectLastValue(underTest.displayId)
-
-            displayRepository.addDisplays(display(id = 2, type = TYPE_EXTERNAL))
-            underTest.onStatusBarTouched(createMotionEventForDisplay(2), STATUS_BAR_WIDTH)
-
-            assertThat(displayId).isEqualTo(2)
-
-            keyguardRepository.setKeyguardShowing(true)
-
-            assertThat(displayId).isEqualTo(Display.DEFAULT_DISPLAY)
-        }
-
-    @Test
-    fun onStatusBarTouched_afterKeyguardHides_goesBackToPreviousDisplay() =
-        testScope.runTest {
-            val underTest = createUnderTest(shadeOnDefaultDisplayWhenLocked = true)
-            val displayId by collectLastValue(underTest.displayId)
-
-            displayRepository.addDisplays(display(id = 2, type = TYPE_EXTERNAL))
-            underTest.onStatusBarTouched(createMotionEventForDisplay(2), STATUS_BAR_WIDTH)
-
-            assertThat(displayId).isEqualTo(2)
-
-            keyguardRepository.setKeyguardShowing(true)
-
-            assertThat(displayId).isEqualTo(Display.DEFAULT_DISPLAY)
-
-            keyguardRepository.setKeyguardShowing(false)
-
-            assertThat(displayId).isEqualTo(2)
-        }
-
-    @Test
     fun onStatusBarTouched_leftSide_intentSetToNotifications() =
         testScope.runTest {
-            val underTest = createUnderTest(shadeOnDefaultDisplayWhenLocked = true)
-
             underTest.onStatusBarTouched(
                 createMotionEventForDisplay(2, STATUS_BAR_WIDTH * 0.1f),
                 STATUS_BAR_WIDTH,
@@ -169,8 +112,6 @@ class StatusBarTouchShadeDisplayPolicyTest : SysuiTestCase() {
     @Test
     fun onStatusBarTouched_rightSide_intentSetToQs() =
         testScope.runTest {
-            val underTest = createUnderTest(shadeOnDefaultDisplayWhenLocked = true)
-
             underTest.onStatusBarTouched(
                 createMotionEventForDisplay(2, STATUS_BAR_WIDTH * 0.95f),
                 STATUS_BAR_WIDTH,
@@ -182,8 +123,6 @@ class StatusBarTouchShadeDisplayPolicyTest : SysuiTestCase() {
     @Test
     fun onStatusBarTouched_nullAfterConsumed() =
         testScope.runTest {
-            val underTest = createUnderTest(shadeOnDefaultDisplayWhenLocked = true)
-
             underTest.onStatusBarTouched(
                 createMotionEventForDisplay(2, STATUS_BAR_WIDTH * 0.1f),
                 STATUS_BAR_WIDTH,
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ShadeDisplayAwareModule.kt b/packages/SystemUI/src/com/android/systemui/shade/ShadeDisplayAwareModule.kt
index 747642097327..f926d39760fe 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadeDisplayAwareModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadeDisplayAwareModule.kt
@@ -36,7 +36,6 @@ import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.res.R
 import com.android.systemui.scene.ui.view.WindowRootView
-import com.android.systemui.shade.data.repository.MutableShadeDisplaysRepository
 import com.android.systemui.shade.data.repository.ShadeDisplaysRepository
 import com.android.systemui.shade.data.repository.ShadeDisplaysRepositoryImpl
 import com.android.systemui.shade.display.ShadeDisplayPolicyModule
@@ -211,15 +210,6 @@ object ShadeDisplayAwareModule {
         return impl
     }
 
-    @SysUISingleton
-    @Provides
-    fun provideMutableShadePositionRepository(
-        impl: ShadeDisplaysRepositoryImpl
-    ): MutableShadeDisplaysRepository {
-        ShadeWindowGoesAround.isUnexpectedlyInLegacyMode()
-        return impl
-    }
-
     @Provides
     @SysUISingleton
     fun provideShadeDialogContextInteractor(
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ShadePrimaryDisplayCommand.kt b/packages/SystemUI/src/com/android/systemui/shade/ShadePrimaryDisplayCommand.kt
index 7bfe40c3d811..173da336c62f 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadePrimaryDisplayCommand.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadePrimaryDisplayCommand.kt
@@ -20,7 +20,7 @@ import android.provider.Settings.Global.DEVELOPMENT_SHADE_DISPLAY_AWARENESS
 import com.android.systemui.CoreStartable
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.display.data.repository.DisplayRepository
-import com.android.systemui.shade.data.repository.MutableShadeDisplaysRepository
+import com.android.systemui.shade.data.repository.ShadeDisplaysRepository
 import com.android.systemui.shade.display.ShadeDisplayPolicy
 import com.android.systemui.statusbar.commandline.Command
 import com.android.systemui.statusbar.commandline.CommandRegistry
@@ -35,7 +35,7 @@ constructor(
     private val globalSettings: GlobalSettings,
     private val commandRegistry: CommandRegistry,
     private val displaysRepository: DisplayRepository,
-    private val positionRepository: MutableShadeDisplaysRepository,
+    private val positionRepository: ShadeDisplaysRepository,
     private val policies: Set<@JvmSuppressWildcards ShadeDisplayPolicy>,
     private val defaultPolicy: ShadeDisplayPolicy,
 ) : Command, CoreStartable {
@@ -103,7 +103,7 @@ constructor(
         }
 
         private fun printPolicies() {
-            val currentPolicyName = positionRepository.policy.value.name
+            val currentPolicyName = positionRepository.currentPolicy.name
             pw.println("Available policies: ")
             policies.forEach {
                 pw.print(" - ${it.name}")
diff --git a/packages/SystemUI/src/com/android/systemui/shade/data/repository/FakeShadeDisplayRepository.kt b/packages/SystemUI/src/com/android/systemui/shade/data/repository/FakeShadeDisplayRepository.kt
index 732d4d1500e7..3513334f2a5c 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/data/repository/FakeShadeDisplayRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/data/repository/FakeShadeDisplayRepository.kt
@@ -17,6 +17,8 @@
 package com.android.systemui.shade.data.repository
 
 import android.view.Display
+import com.android.systemui.shade.display.FakeShadeDisplayPolicy
+import com.android.systemui.shade.display.ShadeDisplayPolicy
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 
@@ -30,7 +32,6 @@ class FakeShadeDisplayRepository : ShadeDisplaysRepository {
     override val displayId: StateFlow<Int>
         get() = _displayId
 
-    fun resetDisplayId() {
-        _displayId.value = Display.DEFAULT_DISPLAY
-    }
+    override val currentPolicy: ShadeDisplayPolicy
+        get() = FakeShadeDisplayPolicy
 }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepository.kt b/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepository.kt
index af48231e0a99..f959f7fe0c31 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepository.kt
@@ -20,14 +20,18 @@ import android.provider.Settings.Global.DEVELOPMENT_SHADE_DISPLAY_AWARENESS
 import android.view.Display
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.keyguard.data.repository.KeyguardRepository
+import com.android.systemui.shade.ShadeOnDefaultDisplayWhenLocked
 import com.android.systemui.shade.display.ShadeDisplayPolicy
 import com.android.systemui.util.settings.GlobalSettings
 import com.android.systemui.util.settings.SettingsProxyExt.observerFlow
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.map
@@ -38,12 +42,8 @@ import kotlinx.coroutines.flow.stateIn
 interface ShadeDisplaysRepository {
     /** ID of the display which currently hosts the shade */
     val displayId: StateFlow<Int>
-}
-
-/** Allows to change the policy that determines in which display the Shade window is visible. */
-interface MutableShadeDisplaysRepository : ShadeDisplaysRepository {
-    /** Updates the policy to select where the shade is visible. */
-    val policy: StateFlow<ShadeDisplayPolicy>
+    /** The current policy set. */
+    val currentPolicy: ShadeDisplayPolicy
 }
 
 /** Keeps the policy and propagates the display id for the shade from it. */
@@ -56,9 +56,11 @@ constructor(
     defaultPolicy: ShadeDisplayPolicy,
     @Background bgScope: CoroutineScope,
     policies: Set<@JvmSuppressWildcards ShadeDisplayPolicy>,
-) : MutableShadeDisplaysRepository {
+    @ShadeOnDefaultDisplayWhenLocked private val shadeOnDefaultDisplayWhenLocked: Boolean,
+    keyguardRepository: KeyguardRepository,
+) : ShadeDisplaysRepository {
 
-    override val policy: StateFlow<ShadeDisplayPolicy> =
+    private val policy: StateFlow<ShadeDisplayPolicy> =
         globalSettings
             .observerFlow(DEVELOPMENT_SHADE_DISPLAY_AWARENESS)
             .onStart { emit(Unit) }
@@ -71,10 +73,32 @@ constructor(
                 return@map defaultPolicy
             }
             .distinctUntilChanged()
-            .stateIn(bgScope, SharingStarted.WhileSubscribed(), defaultPolicy)
+            .stateIn(bgScope, SharingStarted.Eagerly, defaultPolicy)
+
+    private val displayIdFromPolicy: Flow<Int> = policy.flatMapLatest { it.displayId }
+
+    private val keyguardAwareDisplayPolicy: Flow<Int> =
+        if (!shadeOnDefaultDisplayWhenLocked) {
+            displayIdFromPolicy
+        } else {
+            keyguardRepository.isKeyguardShowing.combine(displayIdFromPolicy) {
+                isKeyguardShowing,
+                currentDisplayId ->
+                if (isKeyguardShowing) {
+                    Display.DEFAULT_DISPLAY
+                } else {
+                    currentDisplayId
+                }
+            }
+        }
+
+    override val currentPolicy: ShadeDisplayPolicy
+        get() = policy.value
 
     override val displayId: StateFlow<Int> =
-        policy
-            .flatMapLatest { it.displayId }
-            .stateIn(bgScope, SharingStarted.WhileSubscribed(), Display.DEFAULT_DISPLAY)
+        keyguardAwareDisplayPolicy.stateIn(
+            bgScope,
+            SharingStarted.WhileSubscribed(),
+            Display.DEFAULT_DISPLAY,
+        )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/display/FakeShadeDisplayPolicy.kt b/packages/SystemUI/src/com/android/systemui/shade/display/FakeShadeDisplayPolicy.kt
new file mode 100644
index 000000000000..e010bd6f9880
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/shade/display/FakeShadeDisplayPolicy.kt
@@ -0,0 +1,36 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.shade.display
+
+import android.view.Display
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+
+/** Used only for testing. */
+object FakeShadeDisplayPolicy : ShadeDisplayPolicy {
+    override val name: String
+        get() = "fake_shade_policy"
+
+    override val displayId: StateFlow<Int>
+        get() = _displayId
+
+    private val _displayId = MutableStateFlow(Display.DEFAULT_DISPLAY)
+
+    fun setDisplayId(displayId: Int) {
+        _displayId.value = displayId
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicy.kt b/packages/SystemUI/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicy.kt
index 91020aa7bdb0..b155ada87efd 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicy.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/display/StatusBarTouchShadeDisplayPolicy.kt
@@ -23,8 +23,6 @@ import com.android.app.tracing.coroutines.launchTraced
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.display.data.repository.DisplayRepository
-import com.android.systemui.keyguard.data.repository.KeyguardRepository
-import com.android.systemui.shade.ShadeOnDefaultDisplayWhenLocked
 import com.android.systemui.shade.domain.interactor.NotificationShadeElement
 import com.android.systemui.shade.domain.interactor.QSShadeElement
 import com.android.systemui.shade.domain.interactor.ShadeExpandedStateInteractor.ShadeElement
@@ -38,13 +36,10 @@ import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.collectLatest
-import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.stateIn
 
 /**
  * Moves the shade on the last display that received a status bar touch.
@@ -57,9 +52,7 @@ class StatusBarTouchShadeDisplayPolicy
 @Inject
 constructor(
     displayRepository: DisplayRepository,
-    keyguardRepository: KeyguardRepository,
     @Background private val backgroundScope: CoroutineScope,
-    @ShadeOnDefaultDisplayWhenLocked private val shadeOnDefaultDisplayWhenLocked: Boolean,
     private val shadeInteractor: Lazy<ShadeInteractor>,
     private val qsShadeElement: Lazy<QSShadeElement>,
     private val notificationElement: Lazy<NotificationShadeElement>,
@@ -72,20 +65,7 @@ constructor(
     private var latestIntent = AtomicReference<ShadeElement?>()
     private var timeoutJob: Job? = null
 
-    override val displayId: StateFlow<Int> =
-        if (shadeOnDefaultDisplayWhenLocked) {
-            keyguardRepository.isKeyguardShowing
-                .combine(currentDisplayId) { isKeyguardShowing, currentDisplayId ->
-                    if (isKeyguardShowing) {
-                        Display.DEFAULT_DISPLAY
-                    } else {
-                        currentDisplayId
-                    }
-                }
-                .stateIn(backgroundScope, SharingStarted.WhileSubscribed(), currentDisplayId.value)
-        } else {
-            currentDisplayId
-        }
+    override val displayId: StateFlow<Int> = currentDisplayId
 
     private var removalListener: Job? = null
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepositoryKosmos.kt
index aaef27d257c5..4d718744320d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/ShadeDisplaysRepositoryKosmos.kt
@@ -22,6 +22,7 @@ import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.shade.display.AnyExternalShadeDisplayPolicy
 import com.android.systemui.shade.display.DefaultDisplayShadePolicy
+import com.android.systemui.shade.display.FakeShadeDisplayPolicy
 import com.android.systemui.shade.display.ShadeDisplayPolicy
 import com.android.systemui.shade.display.ShadeExpansionIntent
 import com.android.systemui.shade.display.StatusBarTouchShadeDisplayPolicy
@@ -46,8 +47,6 @@ val Kosmos.statusBarTouchShadeDisplayPolicy: StatusBarTouchShadeDisplayPolicy by
         StatusBarTouchShadeDisplayPolicy(
             displayRepository = displayRepository,
             backgroundScope = testScope.backgroundScope,
-            keyguardRepository = keyguardRepository,
-            shadeOnDefaultDisplayWhenLocked = false,
             shadeInteractor = { shadeInteractor },
             notificationElement = { notificationElement },
             qsShadeElement = { qsElement },
@@ -55,13 +54,15 @@ val Kosmos.statusBarTouchShadeDisplayPolicy: StatusBarTouchShadeDisplayPolicy by
     }
 val Kosmos.shadeExpansionIntent: ShadeExpansionIntent by
     Kosmos.Fixture { statusBarTouchShadeDisplayPolicy }
-val Kosmos.shadeDisplaysRepository: MutableShadeDisplaysRepository by
+val Kosmos.shadeDisplaysRepository: ShadeDisplaysRepository by
     Kosmos.Fixture {
         ShadeDisplaysRepositoryImpl(
             bgScope = testScope.backgroundScope,
             globalSettings = fakeGlobalSettings,
             policies = shadeDisplayPolicies,
             defaultPolicy = defaultShadeDisplayPolicy,
+            shadeOnDefaultDisplayWhenLocked = true,
+            keyguardRepository = keyguardRepository,
         )
     }
 
@@ -71,6 +72,7 @@ val Kosmos.shadeDisplayPolicies: Set<ShadeDisplayPolicy> by
             defaultShadeDisplayPolicy,
             anyExternalShadeDisplayPolicy,
             statusBarTouchShadeDisplayPolicy,
+            FakeShadeDisplayPolicy,
         )
     }