Disable multidisplay in shelf UI

Some bugs need to be ironed out until it can be reenabled safely.

Bug: 339424226
Test: atest TakeScreenshotExecutorTest
Flag: ACONFIG com.android.systemui.screenshot_shelf_ui2 NEXTFOOD
Change-Id: Iec6422567d216ac944be1f4119a5473d7556f21f
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
index 5feac80..c78b63a 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
@@ -8,6 +8,7 @@
 import com.android.app.tracing.coroutines.launch
 import com.android.internal.logging.UiEventLogger
 import com.android.internal.util.ScreenshotRequest
+import com.android.systemui.Flags.screenshotShelfUi2
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.display.data.repository.DisplayRepository
@@ -123,8 +124,8 @@
 
     private suspend fun getDisplaysToScreenshot(requestType: Int): List<Display> {
         val allDisplays = displays.first()
-        return if (requestType == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
-            // If this is a provided image, let's show the UI on the default display only.
+        return if (requestType == TAKE_SCREENSHOT_PROVIDED_IMAGE || screenshotShelfUi2()) {
+            // If this is a provided image or using the shelf UI, just screenshot th default display
             allDisplays.filter { it.displayId == Display.DEFAULT_DISPLAY }
         } else {
             allDisplays.filter { it.type in ALLOWED_DISPLAY_TYPES }
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 0dda41f..e4ced00 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
@@ -3,6 +3,8 @@
 import android.content.ComponentName
 import android.graphics.Bitmap
 import android.net.Uri
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import android.testing.AndroidTestingRunner
 import android.view.Display
 import android.view.Display.TYPE_EXTERNAL
@@ -15,6 +17,7 @@
 import androidx.test.filters.SmallTest
 import com.android.internal.logging.testing.UiEventLoggerFake
 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.display
@@ -77,6 +80,7 @@
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_severalDisplays_callsControllerForEachOne() =
         testScope.runTest {
             val internalDisplay = display(TYPE_INTERNAL, id = 0)
@@ -108,6 +112,32 @@
         }
 
     @Test
+    @EnableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
+    fun executeScreenshots_severalDisplaysShelfUi_justCallsOne() =
+        testScope.runTest {
+            val internalDisplay = display(TYPE_INTERNAL, id = 0)
+            val externalDisplay = display(TYPE_EXTERNAL, id = 1)
+            setDisplays(internalDisplay, externalDisplay)
+            val onSaved = { _: Uri? -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            verify(controllerFactory).create(eq(internalDisplay), any())
+
+            val capturer = ArgumentCaptor<ScreenshotData>()
+
+            verify(controller0).handleScreenshot(capturer.capture(), any(), any())
+            assertThat(capturer.value.displayId).isEqualTo(0)
+
+            assertThat(eventLogger.numLogs()).isEqualTo(1)
+            assertThat(eventLogger.get(0).eventId)
+                .isEqualTo(ScreenshotEvent.SCREENSHOT_REQUESTED_KEY_OTHER.id)
+            assertThat(eventLogger.get(0).packageName).isEqualTo(topComponent.packageName)
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_providedImageType_callsOnlyDefaultDisplayController() =
         testScope.runTest {
             val internalDisplay = display(TYPE_INTERNAL, id = 0)
@@ -139,6 +169,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_onlyVirtualDisplays_noInteractionsWithControllers() =
         testScope.runTest {
             setDisplays(display(TYPE_VIRTUAL, id = 0), display(TYPE_VIRTUAL, id = 1))
@@ -150,6 +181,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_allowedTypes_allCaptured() =
         testScope.runTest {
             whenever(controllerFactory.create(any(), any())).thenReturn(controller0)
@@ -168,6 +200,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_reportsOnFinishedOnlyWhenBothFinished() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -193,6 +226,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_oneFinishesOtherFails_reportFailsOnlyAtTheEnd() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -220,6 +254,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_allDisplaysFail_reportsFail() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -247,6 +282,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun onDestroy_propagatedToControllers() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -259,6 +295,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun removeWindows_propagatedToControllers() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -273,6 +310,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun onCloseSystemDialogsReceived_propagatedToControllers() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -287,6 +325,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun onCloseSystemDialogsReceived_someControllerHavePendingTransitions() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -303,6 +342,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_controllerCalledWithRequestProcessorReturnValue() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0))
@@ -324,6 +364,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_errorFromProcessor_logsScreenshotRequested() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -341,6 +382,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_errorFromProcessor_logsUiError() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -358,6 +400,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_errorFromProcessorOnDefaultDisplay_showsErrorNotification() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -384,6 +427,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_errorFromScreenshotController_reportsRequested() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -404,6 +448,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_errorFromScreenshotController_reportsError() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
@@ -424,6 +469,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_SHELF_UI2)
     fun executeScreenshots_errorFromScreenshotController_showsErrorNotification() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))