Remove duplicated screenshot capture.

It's only being triggered in screenshot controller[s] if it fails in the
policy processor. That first failure should trigger the error case
directly (it's already attached to a runCatching block that shows an
error notification if needed).

Test: Force error, verify that notification appears.
Test: atest PolicyRequestProcessorTest
Flag: EXEMPT duplicated functionality removal
Bug: 371568662
Change-Id: I7954f78e7fbea1e96d388a4ca919926a87fe59e2
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/HeadlessScreenshotHandler.kt b/packages/SystemUI/src/com/android/systemui/screenshot/HeadlessScreenshotHandler.kt
index 7b56688..6d3b9aa 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/HeadlessScreenshotHandler.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/HeadlessScreenshotHandler.kt
@@ -19,7 +19,6 @@
 import android.net.Uri
 import android.os.UserManager
 import android.util.Log
-import android.view.WindowManager
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.res.R
@@ -51,10 +50,6 @@
         finisher: Consumer<Uri?>,
         requestCallback: TakeScreenshotService.RequestCallback,
     ) {
-        if (screenshot.type == WindowManager.TAKE_SCREENSHOT_FULLSCREEN) {
-            screenshot.bitmap = imageCapture.captureDisplay(screenshot.displayId, crop = null)
-        }
-
         if (screenshot.bitmap == null) {
             Log.e(TAG, "handleScreenshot: Screenshot bitmap was null")
             notificationsControllerFactory
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
index 3920d58..63ced5d 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
@@ -255,12 +255,6 @@
         Assert.isMainThread();
 
         mCurrentRequestCallback = requestCallback;
-        Rect bounds = screenshot.getOriginalScreenBounds();
-        if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_FULLSCREEN
-                && screenshot.getBitmap() == null) {
-            bounds = getFullScreenRect();
-            screenshot.setBitmap(mImageCapture.captureDisplay(mDisplay.getDisplayId(), bounds));
-        }
 
         if (screenshot.getBitmap() == null) {
             Log.e(TAG, "handleScreenshot: Screenshot bitmap was null");
@@ -323,6 +317,7 @@
 
         attachWindow();
 
+        Rect bounds = screenshot.getOriginalScreenBounds();
         boolean showFlash;
         if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE) {
             if (bounds != null
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt
index e3fbb60..f5c6052 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt
@@ -35,7 +35,6 @@
 import android.view.Display
 import android.view.ScrollCaptureResponse
 import android.view.ViewRootImpl.ActivityConfigCallback
-import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
 import android.widget.Toast
 import android.window.WindowContext
@@ -162,11 +161,6 @@
         Assert.isMainThread()
         screenshotHandler.resetTimeout()
 
-        if (screenshot.type == TAKE_SCREENSHOT_FULLSCREEN && screenshot.bitmap == null) {
-            val bounds = fullScreenRect
-            screenshot.bitmap = imageCapture.captureDisplay(display.displayId, bounds)
-        }
-
         val currentBitmap = screenshot.bitmap
         if (currentBitmap == null) {
             Log.e(TAG, "handleScreenshot: Screenshot bitmap was null")
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
index a94393b..039143a 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
@@ -126,7 +126,7 @@
         )
     }
 
-    suspend fun replaceWithTaskSnapshot(
+    private suspend fun replaceWithTaskSnapshot(
         original: ScreenshotData,
         componentName: ComponentName?,
         owner: UserHandle,
@@ -134,7 +134,7 @@
         taskBounds: Rect?,
     ): ScreenshotData {
         Log.i(TAG, "Capturing task snapshot: $componentName / $owner")
-        val taskSnapshot = capture.captureTask(taskId)
+        val taskSnapshot = capture.captureTask(taskId) ?: error("Failed to capture task")
         return original.copy(
             type = TAKE_SCREENSHOT_PROVIDED_IMAGE,
             bitmap = taskSnapshot,
@@ -153,13 +153,13 @@
         taskId: Int? = null,
     ): ScreenshotData {
         Log.i(TAG, "Capturing screenshot: $componentName / $owner")
-        val screenshot = captureDisplay(displayId)
+        val screenshot = captureDisplay(displayId) ?: error("Failed to capture screenshot")
         return original.copy(
             type = TAKE_SCREENSHOT_FULLSCREEN,
             bitmap = screenshot,
             userHandle = owner,
             topComponent = componentName,
-            originalScreenBounds = Rect(0, 0, screenshot?.width ?: 0, screenshot?.height ?: 0),
+            originalScreenBounds = Rect(0, 0, screenshot.width, screenshot.height),
             taskId = taskId ?: -1,
         )
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
index e4329a5..0d4cb4c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
@@ -16,6 +16,7 @@
 package com.android.systemui.screenshot.policy
 
 import android.content.ComponentName
+import android.graphics.Bitmap
 import android.graphics.Insets
 import android.graphics.Rect
 import android.os.UserHandle
@@ -27,10 +28,12 @@
 import com.android.systemui.screenshot.ScreenshotData
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.launcherOnly
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
 import com.android.systemui.screenshot.data.repository.DisplayContentRepository
 import com.android.systemui.screenshot.policy.TestUserIds.PERSONAL
 import com.android.systemui.screenshot.policy.TestUserIds.WORK
+import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
@@ -39,14 +42,6 @@
 
 @RunWith(AndroidJUnit4::class)
 class PolicyRequestProcessorTest {
-
-    val imageCapture =
-        object : ImageCapture {
-            override fun captureDisplay(displayId: Int, crop: Rect?) = null
-
-            override suspend fun captureTask(taskId: Int) = null
-        }
-
     /** Tests behavior when no policies are applied */
     @Test
     fun testProcess_defaultOwner_whenNoPolicyApplied() {
@@ -71,7 +66,7 @@
         val requestProcessor =
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
-                imageCapture,
+                createImageCapture(),
                 policies = emptyList(),
                 defaultOwner = UserHandle.of(PERSONAL),
                 defaultComponent = ComponentName("default", "Component"),
@@ -91,6 +86,70 @@
         assertWithMessage("Task ID").that(result.taskId).isEqualTo(TASK_ID)
     }
 
+    @Test
+    fun testProcess_throwsWhenCaptureFails() {
+        val request = ScreenshotData.forTesting()
+
+        /* Create a policy request processor with no capture policies */
+        val requestProcessor =
+            PolicyRequestProcessor(
+                Dispatchers.Unconfined,
+                createImageCapture(display = null),
+                policies = emptyList(),
+                defaultComponent = ComponentName("default", "Component"),
+                displayTasks = DisplayContentRepository { launcherOnly() },
+            )
+
+        val result = runCatching { runBlocking { requestProcessor.process(request) } }
+
+        assertThat(result.isFailure).isTrue()
+    }
+
+    @Test
+    fun testProcess_throwsWhenTaskCaptureFails() {
+        val request = ScreenshotData.forTesting()
+        val fullScreenWork = DisplayContentRepository {
+            singleFullScreen(TaskSpec(taskId = TASK_ID, name = FILES, userId = WORK))
+        }
+
+        val captureTaskPolicy = CapturePolicy {
+            CapturePolicy.PolicyResult.Matched(
+                policy = "",
+                reason = "",
+                parameters =
+                    CaptureParameters(
+                        CaptureType.IsolatedTask(taskId = 0, taskBounds = null),
+                        null,
+                        UserHandle.CURRENT,
+                    ),
+            )
+        }
+
+        /* Create a policy request processor with no capture policies */
+        val requestProcessor =
+            PolicyRequestProcessor(
+                Dispatchers.Unconfined,
+                createImageCapture(task = null),
+                policies = listOf(captureTaskPolicy),
+                defaultComponent = ComponentName("default", "Component"),
+                displayTasks = fullScreenWork,
+            )
+
+        val result = runCatching { runBlocking { requestProcessor.process(request) } }
+
+        assertThat(result.isFailure).isTrue()
+    }
+
+    private fun createImageCapture(
+        display: Bitmap? = Bitmap.createBitmap(100, 100, Bitmap.Config.ARGB_8888),
+        task: Bitmap? = Bitmap.createBitmap(100, 100, Bitmap.Config.ARGB_8888),
+    ) =
+        object : ImageCapture {
+            override fun captureDisplay(displayId: Int, crop: Rect?) = display
+
+            override suspend fun captureTask(taskId: Int) = task
+        }
+
     companion object {
         const val TASK_ID = 1001
     }