More minor cleanups to ScreenshotRequest

* Adds screenshotType into the request instead of relying on the
Message it is sent in (msg.what).

* Adds a constructor which accepts type, source and topComponent

* Adds more @IntDef annotations to `int` constructor args

Bug: 231957192
Test: atest ScreenshotHelperTest
Change-Id: If4bf8ef1b889e4e7b2fabf195a538a13664cce12

4 files changed, 67 insertions, 51 deletions

diff --git a/core/java/com/android/internal/util/ScreenshotHelper.java b/core/java/com/android/internal/util/ScreenshotHelper.java
index b9243ece9158..9474f6fc3252 100644
--- a/core/java/com/android/internal/util/ScreenshotHelper.java
+++ b/core/java/com/android/internal/util/ScreenshotHelper.java
@@ -1,6 +1,7 @@
 package com.android.internal.util;
 
 import static android.content.Intent.ACTION_USER_SWITCHED;
+import static android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -27,7 +28,6 @@ import android.os.Parcelable;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.util.Log;
-import android.view.WindowManager;
 import android.view.WindowManager.ScreenshotSource;
 import android.view.WindowManager.ScreenshotType;
 
@@ -42,10 +42,15 @@ public class ScreenshotHelper {
     public static final int SCREENSHOT_MSG_PROCESS_COMPLETE = 2;
 
     /**
-     * Describes a screenshot request (to make it easier to pass data through to the handler).
+     * Describes a screenshot request.
      */
     public static class ScreenshotRequest implements Parcelable {
+        @ScreenshotType
+        private final int mType;
+
+        @ScreenshotSource
         private final int mSource;
+
         private final Bundle mBitmapBundle;
         private final Rect mBoundsInScreen;
         private final Insets mInsets;
@@ -53,20 +58,27 @@ public class ScreenshotHelper {
         private final int mUserId;
         private final ComponentName mTopComponent;
 
-        @VisibleForTesting
-        public ScreenshotRequest(int source) {
-            mSource = source;
-            mBitmapBundle = null;
-            mBoundsInScreen = null;
-            mInsets = null;
-            mTaskId = -1;
-            mUserId = -1;
-            mTopComponent = null;
+
+        public ScreenshotRequest(@ScreenshotType int type, @ScreenshotSource int source) {
+            this(type, source, /* topComponent */ null);
         }
 
-        @VisibleForTesting
-        public ScreenshotRequest(int source, Bundle bitmapBundle, Rect boundsInScreen,
-                Insets insets, int taskId, int userId, ComponentName topComponent) {
+        public ScreenshotRequest(@ScreenshotType int type, @ScreenshotSource int source,
+                ComponentName topComponent) {
+            this(type,
+                source,
+                /* bitmapBundle*/ null,
+                /* boundsInScreen */ null,
+                /* insets */ null,
+                /* taskId */ -1,
+                /* userId */ -1,
+                topComponent);
+        }
+
+        public ScreenshotRequest(@ScreenshotType int type, @ScreenshotSource int source,
+                Bundle bitmapBundle, Rect boundsInScreen, Insets insets, int taskId, int userId,
+                ComponentName topComponent) {
+            mType = type;
             mSource = source;
             mBitmapBundle = bitmapBundle;
             mBoundsInScreen = boundsInScreen;
@@ -77,6 +89,7 @@ public class ScreenshotHelper {
         }
 
         ScreenshotRequest(Parcel in) {
+            mType = in.readInt();
             mSource = in.readInt();
             if (in.readInt() == 1) {
                 mBitmapBundle = in.readBundle(getClass().getClassLoader());
@@ -96,6 +109,12 @@ public class ScreenshotHelper {
             }
         }
 
+        @ScreenshotType
+        public int getType() {
+            return mType;
+        }
+
+        @ScreenshotSource
         public int getSource() {
             return mSource;
         }
@@ -131,6 +150,7 @@ public class ScreenshotHelper {
 
         @Override
         public void writeToParcel(Parcel dest, int flags) {
+            dest.writeInt(mType);
             dest.writeInt(mSource);
             if (mBitmapBundle == null) {
                 dest.writeInt(0);
@@ -208,8 +228,7 @@ public class ScreenshotHelper {
          * Extracts the Bitmap added to a Bundle with {@link #hardwareBitmapToBundle(Bitmap)} .}
          *
          * <p>This Bitmap contains the HardwareBuffer from the original caller, be careful passing
-         * this
-         * Bitmap on to any other source.
+         * this Bitmap on to any other source.
          *
          * @param bundle containing the bitmap
          * @return a hardware Bitmap
@@ -261,16 +280,16 @@ public class ScreenshotHelper {
      * Added to support reducing unit test duration; the method variant without a timeout argument
      * is recommended for general use.
      *
-     * @param screenshotType The type of screenshot, defined by {@link ScreenshotType}
+     * @param type The type of screenshot, defined by {@link ScreenshotType}
      * @param source The source of the screenshot request, defined by {@link ScreenshotSource}
      * @param handler used to process messages received from the screenshot service
      * @param completionConsumer receives the URI of the captured screenshot, once saved or
      *         null if no screenshot was saved
      */
-    public void takeScreenshot(@ScreenshotType int screenshotType, @ScreenshotSource int source,
+    public void takeScreenshot(@ScreenshotType int type, @ScreenshotSource int source,
             @NonNull Handler handler, @Nullable Consumer<Uri> completionConsumer) {
-        ScreenshotRequest screenshotRequest = new ScreenshotRequest(source);
-        takeScreenshot(screenshotType, handler, screenshotRequest, SCREENSHOT_TIMEOUT_MS,
+        ScreenshotRequest screenshotRequest = new ScreenshotRequest(type, source);
+        takeScreenshot(handler, screenshotRequest, SCREENSHOT_TIMEOUT_MS,
                 completionConsumer);
     }
 
@@ -280,7 +299,7 @@ public class ScreenshotHelper {
      * Added to support reducing unit test duration; the method variant without a timeout argument
      * is recommended for general use.
      *
-     * @param screenshotType The type of screenshot, defined by {@link ScreenshotType}
+     * @param type The type of screenshot, defined by {@link ScreenshotType}
      * @param source The source of the screenshot request, defined by {@link ScreenshotSource}
      * @param handler used to process messages received from the screenshot service
      * @param timeoutMs time limit for processing, intended only for testing
@@ -288,10 +307,10 @@ public class ScreenshotHelper {
      *         null if no screenshot was saved
      */
     @VisibleForTesting
-    public void takeScreenshot(@ScreenshotType int screenshotType, @ScreenshotSource int source,
+    public void takeScreenshot(@ScreenshotType int type, @ScreenshotSource int source,
             @NonNull Handler handler, long timeoutMs, @Nullable Consumer<Uri> completionConsumer) {
-        ScreenshotRequest screenshotRequest = new ScreenshotRequest(source);
-        takeScreenshot(screenshotType, handler, screenshotRequest, timeoutMs, completionConsumer);
+        ScreenshotRequest screenshotRequest = new ScreenshotRequest(type, source);
+        takeScreenshot(handler, screenshotRequest, timeoutMs, completionConsumer);
     }
 
     /**
@@ -312,14 +331,12 @@ public class ScreenshotHelper {
             @NonNull Insets insets, int taskId, int userId, ComponentName topComponent,
             @ScreenshotSource int source, @NonNull Handler handler,
             @Nullable Consumer<Uri> completionConsumer) {
-        ScreenshotRequest screenshotRequest = new ScreenshotRequest(source, screenshotBundle,
-                boundsInScreen, insets, taskId, userId, topComponent);
-        takeScreenshot(WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE, handler, screenshotRequest,
-                SCREENSHOT_TIMEOUT_MS,
-                completionConsumer);
+        ScreenshotRequest screenshotRequest = new ScreenshotRequest(TAKE_SCREENSHOT_PROVIDED_IMAGE,
+                source, screenshotBundle, boundsInScreen, insets, taskId, userId, topComponent);
+        takeScreenshot(handler, screenshotRequest, SCREENSHOT_TIMEOUT_MS, completionConsumer);
     }
 
-    private void takeScreenshot(@ScreenshotType int screenshotType, @NonNull Handler handler,
+    private void takeScreenshot(@NonNull Handler handler,
             ScreenshotRequest screenshotRequest, long timeoutMs,
             @Nullable Consumer<Uri> completionConsumer) {
         synchronized (mScreenshotLock) {
@@ -337,7 +354,7 @@ public class ScreenshotHelper {
                 }
             };
 
-            Message msg = Message.obtain(null, screenshotType, screenshotRequest);
+            Message msg = Message.obtain(null, 0, screenshotRequest);
 
             Handler h = new Handler(handler.getLooper()) {
                 @Override
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
index beb54c85e021..4397d3d85d62 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
@@ -18,7 +18,6 @@ package com.android.systemui.screenshot
 
 import android.net.Uri
 import android.util.Log
-import android.view.WindowManager.ScreenshotType
 import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
 import android.view.WindowManager.TAKE_SCREENSHOT_SELECTED_REGION
@@ -37,13 +36,12 @@ internal class RequestProcessor @Inject constructor(
     private val controller: ScreenshotController,
 ) {
     fun processRequest(
-        @ScreenshotType type: Int,
-        onSavedListener: Consumer<Uri>,
         request: ScreenshotRequest,
+        onSavedListener: Consumer<Uri>,
         callback: RequestCallback
     ) {
 
-        if (type == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
+        if (request.type == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
             val image = HardwareBitmapBundler.bundleToHardwareBitmap(request.bitmapBundle)
 
             controller.handleImageAsScreenshot(
@@ -53,12 +51,12 @@ internal class RequestProcessor @Inject constructor(
             return
         }
 
-        when (type) {
+        when (request.type) {
             TAKE_SCREENSHOT_FULLSCREEN ->
                 controller.takeScreenshotFullscreen(null, onSavedListener, callback)
             TAKE_SCREENSHOT_SELECTED_REGION ->
                 controller.takeScreenshotPartial(null, onSavedListener, callback)
-            else -> Log.w(TAG, "Invalid screenshot option: $type")
+            else -> Log.w(TAG, "Invalid screenshot option: ${request.type}")
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
index f1f0223632b7..7bf3217e5f15 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
@@ -229,11 +229,11 @@ public class TakeScreenshotService extends Service {
 
         if (mFeatureFlags.isEnabled(SCREENSHOT_REQUEST_PROCESSOR)) {
             Log.d(TAG, "handleMessage: Using request processor");
-            mProcessor.processRequest(msg.what, uriConsumer, screenshotRequest, requestCallback);
+            mProcessor.processRequest(screenshotRequest, uriConsumer, requestCallback);
             return true;
         }
 
-        switch (msg.what) {
+        switch (screenshotRequest.getType()) {
             case WindowManager.TAKE_SCREENSHOT_FULLSCREEN:
                 if (DEBUG_SERVICE) {
                     Log.d(TAG, "handleMessage: TAKE_SCREENSHOT_FULLSCREEN");
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/RequestProcessorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/RequestProcessorTest.kt
index 002f23a9ed6d..024d3bd8eb0e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/RequestProcessorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/RequestProcessorTest.kt
@@ -23,8 +23,12 @@ import android.graphics.Insets
 import android.graphics.Rect
 import android.hardware.HardwareBuffer
 import android.net.Uri
-import android.view.WindowManager
-import android.view.WindowManager.ScreenshotSource
+import android.view.WindowManager.ScreenshotSource.SCREENSHOT_KEY_CHORD
+import android.view.WindowManager.ScreenshotSource.SCREENSHOT_OTHER
+import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
+import android.view.WindowManager.TAKE_SCREENSHOT_SELECTED_REGION
+import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
+
 import com.android.internal.util.ScreenshotHelper.HardwareBitmapBundler
 import com.android.internal.util.ScreenshotHelper.ScreenshotRequest
 import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback
@@ -43,13 +47,12 @@ class RequestProcessorTest {
 
     @Test
     fun testFullScreenshot() {
-        val request = ScreenshotRequest(ScreenshotSource.SCREENSHOT_KEY_CHORD)
+        val request = ScreenshotRequest(TAKE_SCREENSHOT_FULLSCREEN, SCREENSHOT_KEY_CHORD)
         val onSavedListener = mock<Consumer<Uri>>()
         val callback = mock<RequestCallback>()
         val processor = RequestProcessor(controller)
 
-        processor.processRequest(WindowManager.TAKE_SCREENSHOT_FULLSCREEN, onSavedListener,
-            request, callback)
+        processor.processRequest(request, onSavedListener, callback)
 
         verify(controller).takeScreenshotFullscreen(/* topComponent */ isNull(),
             eq(onSavedListener), eq(callback))
@@ -57,13 +60,12 @@ class RequestProcessorTest {
 
     @Test
     fun testSelectedRegionScreenshot() {
-        val request = ScreenshotRequest(ScreenshotSource.SCREENSHOT_KEY_CHORD)
+        val request = ScreenshotRequest(TAKE_SCREENSHOT_SELECTED_REGION, SCREENSHOT_KEY_CHORD)
         val onSavedListener = mock<Consumer<Uri>>()
         val callback = mock<RequestCallback>()
         val processor = RequestProcessor(controller)
 
-        processor.processRequest(WindowManager.TAKE_SCREENSHOT_SELECTED_REGION, onSavedListener,
-            request, callback)
+        processor.processRequest(request, onSavedListener, callback)
 
         verify(controller).takeScreenshotPartial(/* topComponent */ isNull(),
             eq(onSavedListener), eq(callback))
@@ -82,14 +84,13 @@ class RequestProcessorTest {
         val bitmap = Bitmap.wrapHardwareBuffer(buffer, ColorSpace.get(ColorSpace.Named.SRGB))!!
         val bitmapBundle = HardwareBitmapBundler.hardwareBitmapToBundle(bitmap)
 
-        val request = ScreenshotRequest(ScreenshotSource.SCREENSHOT_OTHER, bitmapBundle,
-            bounds, Insets.NONE, taskId, userId, topComponent)
+        val request = ScreenshotRequest(TAKE_SCREENSHOT_PROVIDED_IMAGE, SCREENSHOT_OTHER,
+            bitmapBundle, bounds, Insets.NONE, taskId, userId, topComponent)
 
         val onSavedListener = mock<Consumer<Uri>>()
         val callback = mock<RequestCallback>()
 
-        processor.processRequest(WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE, onSavedListener,
-            request, callback)
+        processor.processRequest(request, onSavedListener, callback)
 
         verify(controller).handleImageAsScreenshot(
             bitmapCaptor.capture(), eq(bounds), eq(Insets.NONE), eq(taskId), eq(userId),