Make screenshots use the focused display for some screenshot types

Modify FocusedDisplayRepository to make testing easier.

Bug: 362720389
Test: atest TakeScreenshotExecutorTest
Flag: com.android.systemui.screenshot_multidisplay_focus_change
Change-Id: Ie8326399b23c5bc7aa668981a4bfe38f4c21a926

5 files changed, 119 insertions, 9 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt b/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt
index 373279cec5d1..462e820f68da 100644
--- a/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt
@@ -20,6 +20,8 @@ import com.android.systemui.display.data.repository.DeviceStateRepository
 import com.android.systemui.display.data.repository.DeviceStateRepositoryImpl
 import com.android.systemui.display.data.repository.DisplayRepository
 import com.android.systemui.display.data.repository.DisplayRepositoryImpl
+import com.android.systemui.display.data.repository.FocusedDisplayRepository
+import com.android.systemui.display.data.repository.FocusedDisplayRepositoryImpl
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractorImpl
 import dagger.Binds
@@ -39,4 +41,9 @@ interface DisplayModule {
     fun bindsDeviceStateRepository(
         deviceStateRepository: DeviceStateRepositoryImpl
     ): DeviceStateRepository
+
+    @Binds
+    fun bindsFocusedDisplayRepository(
+        focusedDisplayRepository: FocusedDisplayRepositoryImpl
+    ): FocusedDisplayRepository
 }
diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/FocusedDisplayRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/FocusedDisplayRepository.kt
index 3bfc7a1747ec..6fc08f6bfab0 100644
--- a/packages/SystemUI/src/com/android/systemui/display/data/repository/FocusedDisplayRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/FocusedDisplayRepository.kt
@@ -37,16 +37,21 @@ import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.stateIn
 
 /** Repository tracking display focus. */
+interface FocusedDisplayRepository {
+    /** Provides the currently focused display. */
+    val focusedDisplayId: StateFlow<Int>
+}
+
 @SysUISingleton
 @MainThread
-class FocusedDisplayRepository
+class FocusedDisplayRepositoryImpl
 @Inject
 constructor(
     @Background val backgroundScope: CoroutineScope,
     @Background private val backgroundExecutor: Executor,
     transitions: ShellTransitions,
     @FocusedDisplayRepoLog logBuffer: LogBuffer,
-) {
+) : FocusedDisplayRepository {
     val focusedTask: Flow<Int> =
         conflatedCallbackFlow<Int> {
                 val listener =
@@ -67,7 +72,6 @@ constructor(
                 )
             }
 
-    /** Provides the currently focused display. */
-    val focusedDisplayId: StateFlow<Int>
+    override val focusedDisplayId: StateFlow<Int>
         get() = focusedTask.stateIn(backgroundScope, SharingStarted.Eagerly, DEFAULT_DISPLAY)
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
index ab8a9539b7f2..a7557463b12f 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
@@ -28,6 +28,7 @@ import com.android.systemui.Flags.screenshotMultidisplayFocusChange
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.display.data.repository.DisplayRepository
+import com.android.systemui.display.data.repository.FocusedDisplayRepository
 import com.android.systemui.res.R
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_CAPTURE_FAILED
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_DISMISSED_OTHER
@@ -83,6 +84,7 @@ constructor(
     private val uiEventLogger: UiEventLogger,
     private val screenshotNotificationControllerFactory: ScreenshotNotificationsController.Factory,
     private val headlessScreenshotHandler: HeadlessScreenshotHandler,
+    private val focusedDisplayRepository: FocusedDisplayRepository,
 ) : TakeScreenshotExecutor {
     private val displays = displayRepository.displays
     private var screenshotController: InteractiveScreenshotHandler? = null
@@ -216,14 +218,13 @@ constructor(
                     ?: error("Can't find default display")
 
             // All other invocations use the focused display
-            else -> focusedDisplay()
+            else ->
+                displayRepository.getDisplay(focusedDisplayRepository.focusedDisplayId.value)
+                    ?: displayRepository.getDisplay(Display.DEFAULT_DISPLAY)
+                    ?: error("Can't find default display")
         }
     }
 
-    // TODO(b/367394043): Determine the focused display here.
-    private suspend fun focusedDisplay() =
-        displayRepository.getDisplay(Display.DEFAULT_DISPLAY) ?: error("Can't find default display")
-
     /** Propagates the close system dialog signal to the ScreenshotController. */
     override fun onCloseSystemDialogsReceived() {
         if (screenshotController?.isPendingSharedTransition() == false) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
index 0bea56007a01..27e9f07af168 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
@@ -20,6 +20,7 @@ import com.android.internal.util.ScreenshotRequest
 import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.display.data.repository.FakeDisplayRepository
+import com.android.systemui.display.data.repository.FakeFocusedDisplayRepository
 import com.android.systemui.display.data.repository.display
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.eq
@@ -58,6 +59,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     private val testScope = TestScope(UnconfinedTestDispatcher())
     private val eventLogger = UiEventLoggerFake()
     private val headlessHandler = mock<HeadlessScreenshotHandler>()
+    private val focusedDisplayRepository = FakeFocusedDisplayRepository()
 
     private val screenshotExecutor =
         TakeScreenshotExecutorImpl(
@@ -68,6 +70,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
             eventLogger,
             notificationControllerFactory,
             headlessHandler,
+            focusedDisplayRepository,
         )
 
     @Before
@@ -309,6 +312,59 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
         }
 
     @Test
+    @EnableFlags(Flags.FLAG_SCREENSHOT_MULTIDISPLAY_FOCUS_CHANGE)
+    fun executeScreenshots_keyOther_usesFocusedDisplay() =
+        testScope.runTest {
+            val displayId = 1
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = displayId))
+            val onSaved = { _: Uri? -> }
+            focusedDisplayRepository.emit(displayId)
+
+            screenshotExecutor.executeScreenshots(
+                createScreenshotRequest(
+                    source = WindowManager.ScreenshotSource.SCREENSHOT_KEY_OTHER
+                ),
+                onSaved,
+                callback,
+            )
+
+            val dataCaptor = ArgumentCaptor<ScreenshotData>()
+
+            verify(controller).handleScreenshot(dataCaptor.capture(), any(), any())
+
+            assertThat(dataCaptor.value.displayId).isEqualTo(displayId)
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    @EnableFlags(Flags.FLAG_SCREENSHOT_MULTIDISPLAY_FOCUS_CHANGE)
+    fun executeScreenshots_keyOtherInvalidDisplay_usesDefault() =
+        testScope.runTest {
+            setDisplays(
+                display(TYPE_INTERNAL, id = Display.DEFAULT_DISPLAY),
+                display(TYPE_EXTERNAL, id = 1),
+            )
+            focusedDisplayRepository.emit(5) // invalid display
+            val onSaved = { _: Uri? -> }
+            screenshotExecutor.executeScreenshots(
+                createScreenshotRequest(
+                    source = WindowManager.ScreenshotSource.SCREENSHOT_KEY_OTHER
+                ),
+                onSaved,
+                callback,
+            )
+
+            val dataCaptor = ArgumentCaptor<ScreenshotData>()
+
+            verify(controller).handleScreenshot(dataCaptor.capture(), any(), any())
+
+            assertThat(dataCaptor.value.displayId).isEqualTo(Display.DEFAULT_DISPLAY)
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
     fun onDestroy_propagatedToControllers() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeFocusedDisplayRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeFocusedDisplayRepository.kt
new file mode 100644
index 000000000000..83df5d874ad6
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeFocusedDisplayRepository.kt
@@ -0,0 +1,42 @@
+/*
+ * Copyright (C) 2024 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.display.data.repository
+
+import android.view.Display
+import com.android.systemui.dagger.SysUISingleton
+import dagger.Binds
+import dagger.Module
+import javax.inject.Inject
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+
+@SysUISingleton
+/** Fake [FocusedDisplayRepository] for testing. */
+class FakeFocusedDisplayRepository @Inject constructor() : FocusedDisplayRepository {
+    private val flow = MutableStateFlow<Int>(Display.DEFAULT_DISPLAY)
+
+    override val focusedDisplayId: StateFlow<Int>
+        get() = flow.asStateFlow()
+
+    suspend fun emit(focusedDisplay: Int) = flow.emit(focusedDisplay)
+}
+
+@Module
+interface FakeFocusedDisplayRepositoryModule {
+    @Binds fun bindFake(fake: FakeFocusedDisplayRepository): FocusedDisplayRepository
+}