Refactor pending display logic

Introduces the concept of ignored displays in DisplayRepository, and handles the case of more than one pending display.

The PendingDisplay interface is now both in DisplayRepository and ConnectedDisplayInteractor as we the data layer should only communicate with the domain one, and the domain one with ui one, but Interactor pending displays are essentially just proxies.

The new `ignore()` method removes the pending display from the repository. `disable` has been removed from ConnectedDisplayInteractor as not used for now.

+ use isKeyguardShowing instead of isKeyguardUnlocked. Before, we were showing the dialog also in the lockscreen if the unlock method was just swipe. Now, just having the keyguard visible is enough to prevent the dialog.

Test: MirroringConfirmationDialogTest, ConnectedDisplayInteractorTest, DisplayRepositoryTest
Bug: 298023961
Change-Id: I06477a9aa652696d8db5807414ee708fb4fe41f9

8 files changed, 394 insertions, 127 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
index e7bbf97e3b51..f68078a8a340 100644
--- a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
@@ -22,6 +22,7 @@ import android.hardware.display.DisplayManager.EVENT_FLAG_DISPLAY_ADDED
 import android.hardware.display.DisplayManager.EVENT_FLAG_DISPLAY_CHANGED
 import android.hardware.display.DisplayManager.EVENT_FLAG_DISPLAY_REMOVED
 import android.os.Handler
+import android.os.Trace
 import android.util.Log
 import android.view.Display
 import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
@@ -34,9 +35,14 @@ import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flowOn
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.stateIn
 
 /** Provides a [Flow] of [Display] as returned by [DisplayManager]. */
@@ -49,7 +55,25 @@ interface DisplayRepository {
      *
      * When `null`, it means there is no pending display waiting to be enabled.
      */
-    val pendingDisplayId: Flow<Int?>
+    val pendingDisplay: Flow<PendingDisplay?>
+
+    /** Represents a connected display that has not been enabled yet. */
+    interface PendingDisplay {
+        /** Id of the pending display. */
+        val id: Int
+
+        /** Enables the display, making it available to the system. */
+        suspend fun enable()
+
+        /**
+         * Ignores the pending display. When called, this specific display id doesn't appear as
+         * pending anymore until the display is disconnected and reconnected again.
+         */
+        suspend fun ignore()
+
+        /** Disables the display, making it unavailable to the system. */
+        suspend fun disable()
+    }
 }
 
 @SysUISingleton
@@ -62,7 +86,8 @@ constructor(
     @Background backgroundCoroutineDispatcher: CoroutineDispatcher
 ) : DisplayRepository {
 
-    override val displays: Flow<Set<Display>> =
+    // Displays are enabled only after receiving them in [onDisplayAdded]
+    private val enabledDisplays: StateFlow<Set<Display>> =
         conflatedCallbackFlow {
                 val callback =
                     object : DisplayListener {
@@ -99,27 +124,38 @@ constructor(
             displayManager.displays?.toSet() ?: emptySet()
         }
 
-    override val pendingDisplayId: Flow<Int?> =
+    /** Propagate to the listeners only enabled displays */
+    override val displays: Flow<Set<Display>> = enabledDisplays
+
+    private val enabledDisplayIds: Flow<Set<Int>> =
+        enabledDisplays
+            .map { enabledDisplaysSet -> enabledDisplaysSet.map { it.displayId }.toSet() }
+            .debugLog("enabledDisplayIds")
+
+    private val ignoredDisplayIds = MutableStateFlow<Set<Int>>(emptySet())
+
+    /* keeps connected displays until they are disconnected. */
+    private val connectedDisplayIds: StateFlow<Set<Int>> =
         conflatedCallbackFlow {
                 val callback =
                     object : DisplayConnectionListener {
-                        private val pendingIds = mutableSetOf<Int>()
+                        private val connectedIds = mutableSetOf<Int>()
                         override fun onDisplayConnected(id: Int) {
-                            pendingIds += id
-                            trySend(id)
+                            if (DEBUG) {
+                                Log.d(TAG, "$id connected")
+                            }
+                            connectedIds += id
+                            ignoredDisplayIds.value -= id
+                            trySend(connectedIds.toSet())
                         }
 
                         override fun onDisplayDisconnected(id: Int) {
-                            if (id in pendingIds) {
-                                pendingIds -= id
-                                trySend(null)
-                            } else {
-                                Log.e(
-                                    TAG,
-                                    "onDisplayDisconnected received for unknown display. " +
-                                        "id=$id, knownIds=$pendingIds"
-                                )
+                            connectedIds -= id
+                            if (DEBUG) {
+                                Log.d(TAG, "$id disconnected. Connected ids: $connectedIds")
                             }
+                            ignoredDisplayIds.value -= id
+                            trySend(connectedIds.toSet())
                         }
                     }
                 displayManager.registerDisplayListener(
@@ -130,15 +166,80 @@ constructor(
                 awaitClose { displayManager.unregisterDisplayListener(callback) }
             }
             .distinctUntilChanged()
+            .debugLog("connectedDisplayIds")
             .flowOn(backgroundCoroutineDispatcher)
             .stateIn(
                 applicationScope,
                 started = SharingStarted.WhileSubscribed(),
-                initialValue = null
+                initialValue = emptySet()
             )
 
+    /**
+     * Pending displays are the ones connected, but not enabled and not ignored. A connected display
+     * is ignored after the user makes the decision to use it or not. For now, the initial decision
+     * from the user is final and not reversible.
+     */
+    private val pendingDisplayIds: Flow<Set<Int>> =
+        combine(enabledDisplayIds, connectedDisplayIds, ignoredDisplayIds) {
+                enabledDisplaysIds,
+                connectedDisplayIds,
+                ignoredDisplayIds ->
+                if (DEBUG) {
+                    Log.d(
+                        TAG,
+                        "combining enabled: $enabledDisplaysIds, " +
+                            "connected: $connectedDisplayIds, ignored: $ignoredDisplayIds"
+                    )
+                }
+                connectedDisplayIds - enabledDisplaysIds - ignoredDisplayIds
+            }
+            .debugLog("pendingDisplayIds")
+
+    override val pendingDisplay: Flow<DisplayRepository.PendingDisplay?> =
+        pendingDisplayIds
+            .map { pendingDisplayIds ->
+                val id = pendingDisplayIds.maxOrNull() ?: return@map null
+                object : DisplayRepository.PendingDisplay {
+                    override val id = id
+                    override suspend fun enable() {
+                        traceSection("DisplayRepository#enable($id)") {
+                            displayManager.enableConnectedDisplay(id)
+                        }
+                        // After the display has been enabled, it is automatically ignored.
+                        ignore()
+                    }
+
+                    override suspend fun ignore() {
+                        traceSection("DisplayRepository#ignore($id)") {
+                            ignoredDisplayIds.value += id
+                        }
+                    }
+
+                    override suspend fun disable() {
+                        ignore()
+                        traceSection("DisplayRepository#disable($id)") {
+                            displayManager.disableConnectedDisplay(id)
+                        }
+                    }
+                }
+            }
+            .debugLog("pendingDisplay")
+
+    private fun <T> Flow<T>.debugLog(flowName: String): Flow<T> {
+        return if (DEBUG) {
+            this.onEach {
+                Log.d(TAG, "$flowName: $it")
+                Trace.asyncTraceForTrackEnd(Trace.TRACE_TAG_APP, "$TAG#$flowName", 0)
+                Trace.asyncTraceForTrackBegin(Trace.TRACE_TAG_APP, "$TAG#$flowName", "$it", 0)
+            }
+        } else {
+            this
+        }
+    }
+
     private companion object {
         const val TAG = "DisplayRepository"
+        val DEBUG = Log.isLoggable(TAG, Log.DEBUG)
     }
 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/display/domain/interactor/ConnectedDisplayInteractor.kt b/packages/SystemUI/src/com/android/systemui/display/domain/interactor/ConnectedDisplayInteractor.kt
index ef6fa26420a3..11ed96d5c7eb 100644
--- a/packages/SystemUI/src/com/android/systemui/display/domain/interactor/ConnectedDisplayInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/domain/interactor/ConnectedDisplayInteractor.kt
@@ -16,14 +16,12 @@
 
 package com.android.systemui.display.domain.interactor
 
-import android.hardware.display.DisplayManager
 import android.view.Display
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.display.data.repository.DisplayRepository
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor.PendingDisplay
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor.State
 import com.android.systemui.keyguard.data.repository.KeyguardRepository
-import com.android.systemui.util.traceSection
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.combine
@@ -52,13 +50,18 @@ interface ConnectedDisplayInteractor {
         CONNECTED_SECURE,
     }
 
-    /** Represents a connected display that has not been enabled yet. */
+    /** Represents a connected display that has not been enabled yet for the UI layer. */
     interface PendingDisplay {
         /** Enables the display, making it available to the system. */
-        fun enable()
+        suspend fun enable()
 
-        /** Disables the display, making it unavailable to the system. */
-        fun disable()
+        /**
+         * Ignores the pending display.
+         *
+         * When called, this specific display id doesn't appear as pending anymore until the display
+         * is disconnected and reconnected again.
+         */
+        suspend fun ignore()
     }
 }
 
@@ -66,7 +69,6 @@ interface ConnectedDisplayInteractor {
 class ConnectedDisplayInteractorImpl
 @Inject
 constructor(
-    private val displayManager: DisplayManager,
     keyguardRepository: KeyguardRepository,
     displayRepository: DisplayRepository,
 ) : ConnectedDisplayInteractor {
@@ -92,28 +94,19 @@ constructor(
 
     // Provides the pending display only if the lockscreen is unlocked
     override val pendingDisplay: Flow<PendingDisplay?> =
-        displayRepository.pendingDisplayId.combine(keyguardRepository.isKeyguardUnlocked) {
-            pendingDisplayId,
-            keyguardUnlocked ->
-            if (pendingDisplayId != null && keyguardUnlocked) {
-                pendingDisplayId.toPendingDisplay()
+        displayRepository.pendingDisplay.combine(keyguardRepository.isKeyguardShowing) {
+            repositoryPendingDisplay,
+            keyguardShowing ->
+            if (repositoryPendingDisplay != null && !keyguardShowing) {
+                repositoryPendingDisplay.toInteractorPendingDisplay()
             } else {
                 null
             }
         }
 
-    private fun Int.toPendingDisplay() =
+    private fun DisplayRepository.PendingDisplay.toInteractorPendingDisplay(): PendingDisplay =
         object : PendingDisplay {
-            val id = this@toPendingDisplay
-            override fun enable() {
-                traceSection("DisplayRepository#enable($id)") {
-                    displayManager.enableConnectedDisplay(id)
-                }
-            }
-            override fun disable() {
-                traceSection("DisplayRepository#enable($id)") {
-                    displayManager.disableConnectedDisplay(id)
-                }
-            }
+            override suspend fun enable() = this@toInteractorPendingDisplay.enable()
+            override suspend fun ignore() = this@toInteractorPendingDisplay.ignore()
         }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/display/ui/view/MirroringConfirmationDialog.kt b/packages/SystemUI/src/com/android/systemui/display/ui/view/MirroringConfirmationDialog.kt
index 174c6ff04a7d..ecc9d0ef7810 100644
--- a/packages/SystemUI/src/com/android/systemui/display/ui/view/MirroringConfirmationDialog.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/ui/view/MirroringConfirmationDialog.kt
@@ -24,15 +24,22 @@ import android.view.WindowManager
 import android.widget.TextView
 import com.android.systemui.R
 
-/** Dialog used to decide what to do with a connected display. */
+/**
+ * Dialog used to decide what to do with a connected display.
+ *
+ * [onCancelMirroring] is called **only** if mirroring didn't start, or when the dismiss button is
+ * pressed.
+ */
 class MirroringConfirmationDialog(
     context: Context,
     private val onStartMirroringClickListener: View.OnClickListener,
-    private val onDismissClickListener: View.OnClickListener,
+    private val onCancelMirroring: View.OnClickListener,
 ) : Dialog(context, R.style.Theme_SystemUI_Dialog) {
 
     private lateinit var mirrorButton: TextView
     private lateinit var dismissButton: TextView
+    private var enabledPressed = false
+
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
         window?.apply {
@@ -45,10 +52,15 @@ class MirroringConfirmationDialog(
         mirrorButton =
             requireViewById<TextView>(R.id.enable_display).apply {
                 setOnClickListener(onStartMirroringClickListener)
+                enabledPressed = true
             }
         dismissButton =
-            requireViewById<TextView>(R.id.cancel).apply {
-                setOnClickListener(onDismissClickListener)
+            requireViewById<TextView>(R.id.cancel).apply { setOnClickListener(onCancelMirroring) }
+
+        setOnDismissListener {
+            if (!enabledPressed) {
+                onCancelMirroring.onClick(null)
             }
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/display/ui/viewmodel/ConnectingDisplayViewModel.kt b/packages/SystemUI/src/com/android/systemui/display/ui/viewmodel/ConnectingDisplayViewModel.kt
index ece33b7f6032..86ef439361b0 100644
--- a/packages/SystemUI/src/com/android/systemui/display/ui/viewmodel/ConnectingDisplayViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/ui/viewmodel/ConnectingDisplayViewModel.kt
@@ -19,13 +19,16 @@ import android.app.Dialog
 import android.content.Context
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor.PendingDisplay
 import com.android.systemui.display.ui.view.MirroringConfirmationDialog
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.launch
 
 /**
  * Shows/hides a dialog to allow the user to decide whether to use the external display for
@@ -38,6 +41,7 @@ constructor(
     private val context: Context,
     private val connectedDisplayInteractor: ConnectedDisplayInteractor,
     @Application private val scope: CoroutineScope,
+    @Background private val bgDispatcher: CoroutineDispatcher
 ) {
 
     private var dialog: Dialog? = null
@@ -61,10 +65,13 @@ constructor(
             MirroringConfirmationDialog(
                     context,
                     onStartMirroringClickListener = {
-                        pendingDisplay.enable()
+                        scope.launch(bgDispatcher) { pendingDisplay.enable() }
                         hideDialog()
                     },
-                    onDismissClickListener = { hideDialog() }
+                    onCancelMirroring = {
+                        scope.launch(bgDispatcher) { pendingDisplay.ignore() }
+                        hideDialog()
+                    }
                 )
                 .apply { show() }
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
index db7c003ee545..4bd380e927c7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
@@ -52,6 +52,7 @@ class DisplayRepositoryTest : SysuiTestCase() {
 
     private val displayManager = mock<DisplayManager>()
     private val displayListener = kotlinArgumentCaptor<DisplayManager.DisplayListener>()
+    private val connectedDisplayListener = kotlinArgumentCaptor<DisplayManager.DisplayListener>()
 
     private val testHandler = FakeHandler(Looper.getMainLooper())
     private val testScope = TestScope(UnconfinedTestDispatcher())
@@ -114,7 +115,7 @@ class DisplayRepositoryTest : SysuiTestCase() {
 
             // Let's make sure it has *NOT* been unregistered, as there is still a subscriber.
             setDisplays(1)
-            displayListener.value.onDisplayAdded(1)
+            sendOnDisplayAdded(1)
             assertThat(firstSubscriber?.ids()).containsExactly(1)
         }
 
@@ -127,7 +128,7 @@ class DisplayRepositoryTest : SysuiTestCase() {
             val value by latestDisplayFlowValue()
 
             setDisplays(1)
-            displayListener.value.onDisplayAdded(1)
+            sendOnDisplayAdded(1)
 
             assertThat(value?.ids()).containsExactly(1)
         }
@@ -138,13 +139,13 @@ class DisplayRepositoryTest : SysuiTestCase() {
             val value by latestDisplayFlowValue()
 
             setDisplays(1, 2, 3, 4)
-            displayListener.value.onDisplayAdded(1)
-            displayListener.value.onDisplayAdded(2)
-            displayListener.value.onDisplayAdded(3)
-            displayListener.value.onDisplayAdded(4)
+            sendOnDisplayAdded(1)
+            sendOnDisplayAdded(2)
+            sendOnDisplayAdded(3)
+            sendOnDisplayAdded(4)
 
             setDisplays(1, 2, 3)
-            displayListener.value.onDisplayRemoved(4)
+            sendOnDisplayRemoved(4)
 
             assertThat(value?.ids()).containsExactly(1, 2, 3)
         }
@@ -155,10 +156,10 @@ class DisplayRepositoryTest : SysuiTestCase() {
             val value by latestDisplayFlowValue()
 
             setDisplays(1, 2, 3, 4)
-            displayListener.value.onDisplayAdded(1)
-            displayListener.value.onDisplayAdded(2)
-            displayListener.value.onDisplayAdded(3)
-            displayListener.value.onDisplayAdded(4)
+            sendOnDisplayAdded(1)
+            sendOnDisplayAdded(2)
+            sendOnDisplayAdded(3)
+            sendOnDisplayAdded(4)
 
             displayListener.value.onDisplayChanged(4)
 
@@ -168,22 +169,22 @@ class DisplayRepositoryTest : SysuiTestCase() {
     @Test
     fun onDisplayConnected_pendingDisplayReceived() =
         testScope.runTest {
-            val pendingDisplay by latestPendingDisplayFlowValue()
+            val pendingDisplay by lastPendingDisplay()
 
-            displayListener.value.onDisplayConnected(1)
+            sendOnDisplayConnected(1)
 
-            assertThat(pendingDisplay).isEqualTo(1)
+            assertThat(pendingDisplay!!.id).isEqualTo(1)
         }
 
     @Test
     fun onDisplayDisconnected_pendingDisplayNull() =
         testScope.runTest {
-            val pendingDisplay by latestPendingDisplayFlowValue()
-            displayListener.value.onDisplayConnected(1)
+            val pendingDisplay by lastPendingDisplay()
+            sendOnDisplayConnected(1)
 
             assertThat(pendingDisplay).isNotNull()
 
-            displayListener.value.onDisplayDisconnected(1)
+            sendOnDisplayDisconnected(1)
 
             assertThat(pendingDisplay).isNull()
         }
@@ -191,24 +192,162 @@ class DisplayRepositoryTest : SysuiTestCase() {
     @Test
     fun onDisplayDisconnected_unknownDisplay_doesNotSendNull() =
         testScope.runTest {
-            val pendingDisplay by latestPendingDisplayFlowValue()
-            displayListener.value.onDisplayConnected(1)
+            val pendingDisplay by lastPendingDisplay()
+            sendOnDisplayConnected(1)
 
             assertThat(pendingDisplay).isNotNull()
 
-            displayListener.value.onDisplayDisconnected(2)
+            sendOnDisplayDisconnected(2)
 
             assertThat(pendingDisplay).isNotNull()
         }
 
     @Test
-    fun onDisplayConnected_multipleTimes_sendsOnlyTheLastOne() =
+    fun onDisplayConnected_multipleTimes_sendsOnlyTheMaximum() =
         testScope.runTest {
-            val pendingDisplay by latestPendingDisplayFlowValue()
-            displayListener.value.onDisplayConnected(1)
-            displayListener.value.onDisplayConnected(2)
+            val pendingDisplay by lastPendingDisplay()
 
-            assertThat(pendingDisplay).isEqualTo(2)
+            sendOnDisplayConnected(1)
+            sendOnDisplayConnected(2)
+
+            assertThat(pendingDisplay!!.id).isEqualTo(2)
+        }
+
+    @Test
+    fun onPendingDisplay_enable_displayEnabled() =
+        testScope.runTest {
+            val pendingDisplay by lastPendingDisplay()
+
+            sendOnDisplayConnected(1)
+            pendingDisplay!!.enable()
+
+            verify(displayManager).enableConnectedDisplay(eq(1))
+        }
+
+    @Test
+    fun onPendingDisplay_enableBySysui_disabledBySomeoneElse_pendingDisplayStillIgnored() =
+        testScope.runTest {
+            val pendingDisplay by lastPendingDisplay()
+
+            sendOnDisplayConnected(1)
+            pendingDisplay!!.enable()
+            // to mock the display being really enabled:
+            sendOnDisplayAdded(1)
+
+            // Simulate the display being disabled by someone else. Now, sysui will have it in the
+            // "pending displays" list again, but it should be ignored.
+            sendOnDisplayRemoved(1)
+
+            assertThat(pendingDisplay).isNull()
+        }
+
+    @Test
+    fun onPendingDisplay_ignoredBySysui_enabledDisabledBySomeoneElse_pendingDisplayStillIgnored() =
+        testScope.runTest {
+            val pendingDisplay by lastPendingDisplay()
+
+            sendOnDisplayConnected(1)
+            pendingDisplay!!.ignore()
+
+            // to mock the display being enabled and disabled by someone else:
+            sendOnDisplayAdded(1)
+            sendOnDisplayRemoved(1)
+
+            // Sysui already decided to ignore it, so the pending display should be null.
+            assertThat(pendingDisplay).isNull()
+        }
+
+    @Test
+    fun onPendingDisplay_disable_displayDisabled() =
+        testScope.runTest {
+            val pendingDisplay by lastPendingDisplay()
+
+            sendOnDisplayConnected(1)
+            pendingDisplay!!.disable()
+
+            verify(displayManager).disableConnectedDisplay(eq(1))
+        }
+
+    @Test
+    fun onPendingDisplay_ignore_pendingDisplayNull() =
+        testScope.runTest {
+            val pendingDisplay by lastPendingDisplay()
+            sendOnDisplayConnected(1)
+
+            pendingDisplay!!.ignore()
+
+            assertThat(pendingDisplay).isNull()
+            verify(displayManager, never()).disableConnectedDisplay(eq(1))
+            verify(displayManager, never()).enableConnectedDisplay(eq(1))
+        }
+
+    @Test
+    fun onPendingDisplay_enabled_pendingDisplayNull() =
+        testScope.runTest {
+            val pendingDisplay by lastPendingDisplay()
+
+            sendOnDisplayConnected(1)
+            assertThat(pendingDisplay).isNotNull()
+
+            setDisplays(1)
+            sendOnDisplayAdded(1)
+
+            assertThat(pendingDisplay).isNull()
+        }
+
+    @Test
+    fun onPendingDisplay_multipleConnected_oneEnabled_pendingDisplayNotNull() =
+        testScope.runTest {
+            val pendingDisplay by lastPendingDisplay()
+
+            sendOnDisplayConnected(1)
+            sendOnDisplayConnected(2)
+
+            assertThat(pendingDisplay).isNotNull()
+
+            setDisplays(1)
+            sendOnDisplayAdded(1)
+
+            assertThat(pendingDisplay).isNotNull()
+            assertThat(pendingDisplay!!.id).isEqualTo(2)
+
+            setDisplays(1, 2)
+            sendOnDisplayAdded(2)
+
+            assertThat(pendingDisplay).isNull()
+        }
+
+    @Test
+    fun pendingDisplay_connectedDisconnectedAndReconnected_expectedPendingDisplayState() =
+        testScope.runTest {
+            val pendingDisplay by lastPendingDisplay()
+
+            // Plug the cable
+            sendOnDisplayConnected(1)
+
+            // Enable it
+            assertThat(pendingDisplay).isNotNull()
+            pendingDisplay!!.enable()
+
+            // Enabled
+            verify(displayManager).enableConnectedDisplay(1)
+            setDisplays(1)
+            sendOnDisplayAdded(1)
+
+            // No more pending displays
+            assertThat(pendingDisplay).isNull()
+
+            // Let's disconnect the cable
+            setDisplays()
+            sendOnDisplayRemoved(1)
+            sendOnDisplayDisconnected(1)
+
+            assertThat(pendingDisplay).isNull()
+
+            // Let's reconnect it
+            sendOnDisplayConnected(1)
+
+            assertThat(pendingDisplay).isNotNull()
         }
 
     private fun Iterable<Display>.ids(): List<Int> = map { it.displayId }
@@ -216,28 +355,47 @@ class DisplayRepositoryTest : SysuiTestCase() {
     // Wrapper to capture the displayListener.
     private fun TestScope.latestDisplayFlowValue(): FlowValue<Set<Display>?> {
         val flowValue = collectLastValue(displayRepository.displays)
+        captureAddedRemovedListener()
+        return flowValue
+    }
+
+    private fun TestScope.lastPendingDisplay(): FlowValue<DisplayRepository.PendingDisplay?> {
+        val flowValue = collectLastValue(displayRepository.pendingDisplay)
+        captureAddedRemovedListener()
         verify(displayManager)
             .registerDisplayListener(
-                displayListener.capture(),
+                connectedDisplayListener.capture(),
                 eq(testHandler),
-                eq(
-                    DisplayManager.EVENT_FLAG_DISPLAY_ADDED or
-                        DisplayManager.EVENT_FLAG_DISPLAY_CHANGED or
-                        DisplayManager.EVENT_FLAG_DISPLAY_REMOVED
-                )
+                eq(DisplayManager.EVENT_FLAG_DISPLAY_CONNECTION_CHANGED)
             )
         return flowValue
     }
 
-    private fun TestScope.latestPendingDisplayFlowValue(): FlowValue<Int?> {
-        val flowValue = collectLastValue(displayRepository.pendingDisplayId)
+    private fun captureAddedRemovedListener() {
         verify(displayManager)
             .registerDisplayListener(
                 displayListener.capture(),
                 eq(testHandler),
-                eq(DisplayManager.EVENT_FLAG_DISPLAY_CONNECTION_CHANGED)
+                eq(
+                    DisplayManager.EVENT_FLAG_DISPLAY_ADDED or
+                        DisplayManager.EVENT_FLAG_DISPLAY_CHANGED or
+                        DisplayManager.EVENT_FLAG_DISPLAY_REMOVED
+                )
             )
-        return flowValue
+    }
+    private fun sendOnDisplayAdded(id: Int) {
+        displayListener.value.onDisplayAdded(id)
+    }
+    private fun sendOnDisplayRemoved(id: Int) {
+        displayListener.value.onDisplayRemoved(id)
+    }
+
+    private fun sendOnDisplayDisconnected(id: Int) {
+        connectedDisplayListener.value.onDisplayDisconnected(id)
+    }
+
+    private fun sendOnDisplayConnected(id: Int) {
+        connectedDisplayListener.value.onDisplayConnected(id)
     }
 
     private fun setDisplays(displays: List<Display>) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/display/domain/interactor/ConnectedDisplayInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/display/domain/interactor/ConnectedDisplayInteractorTest.kt
index 50617a16ce0b..26ee09412687 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/display/domain/interactor/ConnectedDisplayInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/display/domain/interactor/ConnectedDisplayInteractorTest.kt
@@ -16,7 +16,6 @@
 
 package com.android.systemui.display.domain.interactor
 
-import android.hardware.display.DisplayManager
 import android.testing.AndroidTestingRunner
 import android.testing.TestableLooper
 import android.view.Display
@@ -27,12 +26,11 @@ import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.FlowValue
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.display.data.repository.FakeDisplayRepository
+import com.android.systemui.display.data.repository.createPendingDisplay
 import com.android.systemui.display.data.repository.display
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor.PendingDisplay
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor.State
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
-import com.android.systemui.util.mockito.eq
-import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.TestScope
@@ -41,7 +39,6 @@ import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mockito
 
 @RunWith(AndroidTestingRunner::class)
 @TestableLooper.RunWithLooper
@@ -49,20 +46,15 @@ import org.mockito.Mockito
 @SmallTest
 class ConnectedDisplayInteractorTest : SysuiTestCase() {
 
-    private val displayManager = mock<DisplayManager>()
     private val fakeDisplayRepository = FakeDisplayRepository()
     private val fakeKeyguardRepository = FakeKeyguardRepository()
     private val connectedDisplayStateProvider: ConnectedDisplayInteractor =
-        ConnectedDisplayInteractorImpl(
-            displayManager,
-            fakeKeyguardRepository,
-            fakeDisplayRepository
-        )
+        ConnectedDisplayInteractorImpl(fakeKeyguardRepository, fakeDisplayRepository)
     private val testScope = TestScope(UnconfinedTestDispatcher())
 
     @Before
     fun setup() {
-        fakeKeyguardRepository.setKeyguardUnlocked(true)
+        fakeKeyguardRepository.setKeyguardShowing(false)
     }
 
     @Test
@@ -148,7 +140,7 @@ class ConnectedDisplayInteractorTest : SysuiTestCase() {
     fun pendingDisplay_propagated() =
         testScope.runTest {
             val value by lastPendingDisplay()
-            val pendingDisplayId = 4
+            val pendingDisplayId = createPendingDisplay()
 
             fakeDisplayRepository.emit(pendingDisplayId)
 
@@ -156,51 +148,29 @@ class ConnectedDisplayInteractorTest : SysuiTestCase() {
         }
 
     @Test
-    fun onPendingDisplay_enable_displayEnabled() =
+    fun onPendingDisplay_keyguardShowing_returnsPendingDisplay() =
         testScope.runTest {
+            fakeKeyguardRepository.setKeyguardShowing(true)
             val pendingDisplay by lastPendingDisplay()
 
-            fakeDisplayRepository.emit(1)
-            pendingDisplay!!.enable()
-
-            Mockito.verify(displayManager).enableConnectedDisplay(eq(1))
-        }
-
-    @Test
-    fun onPendingDisplay_disable_displayDisabled() =
-        testScope.runTest {
-            val pendingDisplay by lastPendingDisplay()
-
-            fakeDisplayRepository.emit(1)
-            pendingDisplay!!.disable()
-
-            Mockito.verify(displayManager).disableConnectedDisplay(eq(1))
-        }
-
-    @Test
-    fun onPendingDisplay_keyguardUnlocked_returnsPendingDisplay() =
-        testScope.runTest {
-            fakeKeyguardRepository.setKeyguardUnlocked(false)
-            val pendingDisplay by lastPendingDisplay()
-
-            fakeDisplayRepository.emit(1)
+            fakeDisplayRepository.emit(createPendingDisplay())
             assertThat(pendingDisplay).isNull()
 
-            fakeKeyguardRepository.setKeyguardUnlocked(true)
+            fakeKeyguardRepository.setKeyguardShowing(false)
 
             assertThat(pendingDisplay).isNotNull()
         }
 
     @Test
-    fun onPendingDisplay_keyguardLocked_returnsNull() =
+    fun onPendingDisplay_keyguardShowing_returnsNull() =
         testScope.runTest {
-            fakeKeyguardRepository.setKeyguardUnlocked(true)
+            fakeKeyguardRepository.setKeyguardShowing(false)
             val pendingDisplay by lastPendingDisplay()
 
-            fakeDisplayRepository.emit(1)
+            fakeDisplayRepository.emit(createPendingDisplay())
             assertThat(pendingDisplay).isNotNull()
 
-            fakeKeyguardRepository.setKeyguardUnlocked(false)
+            fakeKeyguardRepository.setKeyguardShowing(true)
 
             assertThat(pendingDisplay).isNull()
         }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/display/ui/view/MirroringConfirmationDialogTest.kt b/packages/SystemUI/tests/src/com/android/systemui/display/ui/view/MirroringConfirmationDialogTest.kt
index 705964736a49..46f758216aae 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/display/ui/view/MirroringConfirmationDialogTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/display/ui/view/MirroringConfirmationDialogTest.kt
@@ -68,6 +68,28 @@ class MirroringConfirmationDialogTest : SysuiTestCase() {
         verify(onStartMirroringCallback, never()).onClick(any())
     }
 
+    @Test
+    fun onCancel_afterEnablingMirroring_cancelCallbackNotCalled() {
+        dialog.show()
+        dialog.requireViewById<View>(R.id.enable_display).callOnClick()
+
+        dialog.cancel()
+
+        verify(onCancelCallback, never()).onClick(any())
+        verify(onStartMirroringCallback).onClick(any())
+    }
+
+    @Test
+    fun onDismiss_afterEnablingMirroring_cancelCallbackNotCalled() {
+        dialog.show()
+        dialog.requireViewById<View>(R.id.enable_display).callOnClick()
+
+        dialog.dismiss()
+
+        verify(onCancelCallback, never()).onClick(any())
+        verify(onStartMirroringCallback).onClick(any())
+    }
+
     @After
     fun teardown() {
         if (::dialog.isInitialized) {
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
index 6f51d1b60822..2ac625d68bfe 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
@@ -30,20 +30,24 @@ fun display(type: Int, flags: Int = 0, id: Int = 0): Display {
     }
 }
 
+/** Creates a mock [DisplayRepository.PendingDisplay]. */
+fun createPendingDisplay(id: Int = 0): DisplayRepository.PendingDisplay =
+    mock<DisplayRepository.PendingDisplay> { whenever(this.id).thenReturn(id) }
+
 /** Fake [DisplayRepository] implementation for testing. */
 class FakeDisplayRepository : DisplayRepository {
     private val flow = MutableSharedFlow<Set<Display>>()
-    private val pendingDisplayFlow = MutableSharedFlow<Int?>()
+    private val pendingDisplayFlow = MutableSharedFlow<DisplayRepository.PendingDisplay?>()
 
     /** Emits [value] as [displays] flow value. */
     suspend fun emit(value: Set<Display>) = flow.emit(value)
 
-    /** Emits [value] as [pendingDisplayId] flow value. */
-    suspend fun emit(value: Int?) = pendingDisplayFlow.emit(value)
+    /** Emits [value] as [pendingDisplay] flow value. */
+    suspend fun emit(value: DisplayRepository.PendingDisplay?) = pendingDisplayFlow.emit(value)
 
     override val displays: Flow<Set<Display>>
         get() = flow
 
-    override val pendingDisplayId: Flow<Int?>
+    override val pendingDisplay: Flow<DisplayRepository.PendingDisplay?>
         get() = pendingDisplayFlow
 }