Provide game task bitmap to TakeScreenshotService as screenshot

* Move the BitmapUtil to com.android.internal package
* Remove bitmap result from GameSession#takeScreenshot API

Bug: 219992742
Test: atest GameServiceProviderInstanceImplTest
Change-Id: I4bf29d623f781434ec7ffe4443e658880c31e619
Move the BitmapUtil to com.android.internal package

10 files changed, 163 insertions, 159 deletions

diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 15666017f431..774d7c3785fe 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -11402,7 +11402,7 @@ package android.service.games {
 
   public static interface GameSession.ScreenshotCallback {
     method public void onFailure(int);
-    method public void onSuccess(@NonNull android.graphics.Bitmap);
+    method public void onSuccess();
     field public static final int ERROR_TAKE_SCREENSHOT_INTERNAL_ERROR = 0; // 0x0
   }
 
diff --git a/core/java/android/service/games/GameScreenshotResult.java b/core/java/android/service/games/GameScreenshotResult.java
index ae76e08c7971..5490aef0e225 100644
--- a/core/java/android/service/games/GameScreenshotResult.java
+++ b/core/java/android/service/games/GameScreenshotResult.java
@@ -18,8 +18,6 @@ package android.service.games;
 
 import android.annotation.IntDef;
 import android.annotation.NonNull;
-import android.annotation.Nullable;
-import android.graphics.Bitmap;
 import android.os.Parcel;
 import android.os.Parcelable;
 
@@ -30,9 +28,7 @@ import java.util.Objects;
 /**
  * Result object for calls to {@link IGameSessionController#takeScreenshot}.
  *
- * It includes a status (see {@link #getStatus}) and, if the status is
- * {@link #GAME_SCREENSHOT_SUCCESS} an {@link android.graphics.Bitmap} result (see {@link
- * #getBitmap}).
+ * It includes a status only (see {@link #getStatus}).
  *
  * @hide
  */
@@ -54,8 +50,7 @@ public final class GameScreenshotResult implements Parcelable {
 
     /**
      * Indicates that the result of a call to {@link IGameSessionController#takeScreenshot} was
-     * successful and an {@link android.graphics.Bitmap} result should be available by calling
-     * {@link #getBitmap}.
+     * successful.
      *
      * @hide
      */
@@ -81,9 +76,7 @@ public final class GameScreenshotResult implements Parcelable {
             new Parcelable.Creator<GameScreenshotResult>() {
                 @Override
                 public GameScreenshotResult createFromParcel(Parcel source) {
-                    return new GameScreenshotResult(
-                            source.readInt(),
-                            source.readParcelable(null, Bitmap.class));
+                    return new GameScreenshotResult(source.readInt());
                 }
 
                 @Override
@@ -95,14 +88,11 @@ public final class GameScreenshotResult implements Parcelable {
     @GameScreenshotStatus
     private final int mStatus;
 
-    @Nullable
-    private final Bitmap mBitmap;
-
     /**
-     * Creates a successful {@link GameScreenshotResult} with the provided bitmap.
+     * Creates a successful {@link GameScreenshotResult}.
      */
-    public static GameScreenshotResult createSuccessResult(@NonNull Bitmap bitmap) {
-        return new GameScreenshotResult(GAME_SCREENSHOT_SUCCESS, bitmap);
+    public static GameScreenshotResult createSuccessResult() {
+        return new GameScreenshotResult(GAME_SCREENSHOT_SUCCESS);
     }
 
     /**
@@ -110,12 +100,11 @@ public final class GameScreenshotResult implements Parcelable {
      * {@link #GAME_SCREENSHOT_ERROR_INTERNAL_ERROR} status.
      */
     public static GameScreenshotResult createInternalErrorResult() {
-        return new GameScreenshotResult(GAME_SCREENSHOT_ERROR_INTERNAL_ERROR, null);
+        return new GameScreenshotResult(GAME_SCREENSHOT_ERROR_INTERNAL_ERROR);
     }
 
-    private GameScreenshotResult(@GameScreenshotStatus int status, @Nullable Bitmap bitmap) {
+    private GameScreenshotResult(@GameScreenshotStatus int status) {
         this.mStatus = status;
-        this.mBitmap = bitmap;
     }
 
     @Override
@@ -126,7 +115,6 @@ public final class GameScreenshotResult implements Parcelable {
     @Override
     public void writeToParcel(@NonNull Parcel dest, int flags) {
         dest.writeInt(mStatus);
-        dest.writeParcelable(mBitmap, flags);
     }
 
     @GameScreenshotStatus
@@ -134,29 +122,12 @@ public final class GameScreenshotResult implements Parcelable {
         return mStatus;
     }
 
-    /**
-     * Gets the {@link Bitmap} result from a successful screenshot attempt.
-     *
-     * @return The bitmap.
-     * @throws IllegalStateException if this method is called when {@link #getStatus} does not
-     *                               return {@link #GAME_SCREENSHOT_SUCCESS}.
-     */
-    @NonNull
-    public Bitmap getBitmap() {
-        if (mBitmap == null) {
-            throw new IllegalStateException("Bitmap not available for failed screenshot result");
-        }
-        return mBitmap;
-    }
-
     @Override
     public String toString() {
         return "GameScreenshotResult{"
                 + "mStatus="
                 + mStatus
-                + ", has bitmap='"
-                + mBitmap != null ? "yes" : "no"
-                + "\'}";
+                + "}";
     }
 
     @Override
@@ -170,12 +141,11 @@ public final class GameScreenshotResult implements Parcelable {
         }
 
         GameScreenshotResult that = (GameScreenshotResult) o;
-        return mStatus == that.mStatus
-                && Objects.equals(mBitmap, that.mBitmap);
+        return mStatus == that.mStatus;
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(mStatus, mBitmap);
+        return Objects.hash(mStatus);
     }
 }
diff --git a/core/java/android/service/games/GameSession.java b/core/java/android/service/games/GameSession.java
index 468e087c941b..1548f5224572 100644
--- a/core/java/android/service/games/GameSession.java
+++ b/core/java/android/service/games/GameSession.java
@@ -29,7 +29,6 @@ import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.res.Configuration;
-import android.graphics.Bitmap;
 import android.graphics.Rect;
 import android.os.Binder;
 import android.os.Bundle;
@@ -367,7 +366,7 @@ public abstract class GameSession {
     }
 
     /**
-     * Interface for returning screenshot outcome from calls to {@link #takeScreenshot}.
+     * Interface for handling result of {@link #takeScreenshot}.
      */
     public interface ScreenshotCallback {
 
@@ -402,18 +401,16 @@ public abstract class GameSession {
 
         /**
          * Called when taking the screenshot succeeded.
-         *
-         * @param bitmap The screenshot.
          */
-        void onSuccess(@NonNull Bitmap bitmap);
+        void onSuccess();
     }
 
     /**
      * Takes a screenshot of the associated game. For this call to succeed, the device screen
      * must be turned on and the game task must be visible.
      *
-     * If the callback is called with {@link ScreenshotCallback#onSuccess}, the provided {@link
-     * Bitmap} may be used.
+     * If the callback is called with {@link ScreenshotCallback#onSuccess}, the screenshot is
+     * taken successfully.
      *
      * If the callback is called with {@link ScreenshotCallback#onFailure}, the provided status
      * code should be checked.
@@ -460,7 +457,7 @@ public abstract class GameSession {
         @GameScreenshotResult.GameScreenshotStatus int status = result.getStatus();
         switch (status) {
             case GameScreenshotResult.GAME_SCREENSHOT_SUCCESS:
-                callback.onSuccess(result.getBitmap());
+                callback.onSuccess();
                 break;
             case GameScreenshotResult.GAME_SCREENSHOT_ERROR_INTERNAL_ERROR:
                 Slog.w(TAG, "Error taking screenshot");
diff --git a/core/java/com/android/internal/util/ScreenshotHelper.java b/core/java/com/android/internal/util/ScreenshotHelper.java
index d3c3917cd791..f4f438b1f601 100644
--- a/core/java/com/android/internal/util/ScreenshotHelper.java
+++ b/core/java/com/android/internal/util/ScreenshotHelper.java
@@ -11,8 +11,12 @@ import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.ServiceConnection;
+import android.graphics.Bitmap;
+import android.graphics.ColorSpace;
 import android.graphics.Insets;
+import android.graphics.ParcelableColorSpace;
 import android.graphics.Rect;
+import android.hardware.HardwareBuffer;
 import android.net.Uri;
 import android.os.Bundle;
 import android.os.Handler;
@@ -26,6 +30,7 @@ import android.os.UserHandle;
 import android.util.Log;
 import android.view.WindowManager;
 
+import java.util.Objects;
 import java.util.function.Consumer;
 
 public class ScreenshotHelper {
@@ -154,6 +159,72 @@ public class ScreenshotHelper {
                 };
     }
 
+    /**
+     * Bundler used to convert between a hardware bitmap and a bundle without copying the internal
+     * content. This is expected to be used together with {@link #provideScreenshot} to handle a
+     * hardware bitmap as a screenshot.
+     */
+    public static final class HardwareBitmapBundler {
+        private static final String KEY_BUFFER = "bitmap_util_buffer";
+        private static final String KEY_COLOR_SPACE = "bitmap_util_color_space";
+
+        private HardwareBitmapBundler() {
+        }
+
+        /**
+         * Creates a Bundle that represents the given Bitmap.
+         * <p>The Bundle will contain a wrapped version of the Bitmaps HardwareBuffer, so will avoid
+         * copies when passing across processes, only pass to processes you trust.
+         *
+         * <p>Returns a new Bundle rather than modifying an exiting one to avoid key collisions, the
+         * returned Bundle should be treated as a standalone object.
+         *
+         * @param bitmap to convert to bundle
+         * @return a Bundle representing the bitmap, should only be parsed by
+         * {@link #bundleToHardwareBitmap(Bundle)}
+         */
+        public static Bundle hardwareBitmapToBundle(Bitmap bitmap) {
+            if (bitmap.getConfig() != Bitmap.Config.HARDWARE) {
+                throw new IllegalArgumentException(
+                        "Passed bitmap must have hardware config, found: " + bitmap.getConfig());
+            }
+
+            // Bitmap assumes SRGB for null color space
+            ParcelableColorSpace colorSpace =
+                    bitmap.getColorSpace() == null
+                            ? new ParcelableColorSpace(ColorSpace.get(ColorSpace.Named.SRGB))
+                            : new ParcelableColorSpace(bitmap.getColorSpace());
+
+            Bundle bundle = new Bundle();
+            bundle.putParcelable(KEY_BUFFER, bitmap.getHardwareBuffer());
+            bundle.putParcelable(KEY_COLOR_SPACE, colorSpace);
+
+            return bundle;
+        }
+
+        /**
+         * 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.
+         *
+         * @param bundle containing the bitmap
+         * @return a hardware Bitmap
+         */
+        public static Bitmap bundleToHardwareBitmap(Bundle bundle) {
+            if (!bundle.containsKey(KEY_BUFFER) || !bundle.containsKey(KEY_COLOR_SPACE)) {
+                throw new IllegalArgumentException("Bundle does not contain a hardware bitmap");
+            }
+
+            HardwareBuffer buffer = bundle.getParcelable(KEY_BUFFER);
+            ParcelableColorSpace colorSpace = bundle.getParcelable(KEY_COLOR_SPACE);
+
+            return Bitmap.wrapHardwareBuffer(Objects.requireNonNull(buffer),
+                    colorSpace.getColorSpace());
+        }
+    }
+
     private static final String TAG = "ScreenshotHelper";
 
     // Time until we give up on the screenshot & show an error instead.
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/recents/utilities/BitmapUtil.java b/packages/SystemUI/shared/src/com/android/systemui/shared/recents/utilities/BitmapUtil.java
deleted file mode 100644
index 325bcfc622d7..000000000000
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/recents/utilities/BitmapUtil.java
+++ /dev/null
@@ -1,87 +0,0 @@
-/*
- * Copyright (C) 2020 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.shared.recents.utilities;
-
-import android.graphics.Bitmap;
-import android.graphics.ColorSpace;
-import android.graphics.ParcelableColorSpace;
-import android.hardware.HardwareBuffer;
-import android.os.Bundle;
-
-import java.util.Objects;
-
-/**
- * Utils for working with Bitmaps.
- */
-public final class BitmapUtil {
-    private static final String KEY_BUFFER = "bitmap_util_buffer";
-    private static final String KEY_COLOR_SPACE = "bitmap_util_color_space";
-
-    private BitmapUtil(){ }
-
-    /**
-     * Creates a Bundle that represents the given Bitmap.
-     * <p>The Bundle will contain a wrapped version of the Bitmaps HardwareBuffer, so will avoid
-     * copies when passing across processes, only pass to processes you trust.
-     *
-     * <p>Returns a new Bundle rather than modifying an exiting one to avoid key collisions, the
-     * returned Bundle should be treated as a standalone object.
-     *
-     * @param bitmap to convert to bundle
-     * @return a Bundle representing the bitmap, should only be parsed by
-     *         {@link #bundleToHardwareBitmap(Bundle)}
-     */
-    public static Bundle hardwareBitmapToBundle(Bitmap bitmap) {
-        if (bitmap.getConfig() != Bitmap.Config.HARDWARE) {
-            throw new IllegalArgumentException(
-                    "Passed bitmap must have hardware config, found: " + bitmap.getConfig());
-        }
-
-        // Bitmap assumes SRGB for null color space
-        ParcelableColorSpace colorSpace =
-                bitmap.getColorSpace() == null
-                        ? new ParcelableColorSpace(ColorSpace.get(ColorSpace.Named.SRGB))
-                        : new ParcelableColorSpace(bitmap.getColorSpace());
-
-        Bundle bundle = new Bundle();
-        bundle.putParcelable(KEY_BUFFER, bitmap.getHardwareBuffer());
-        bundle.putParcelable(KEY_COLOR_SPACE, colorSpace);
-
-        return bundle;
-    }
-
-    /**
-     * 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.
-     *
-     * @param bundle containing the bitmap
-     * @return a hardware Bitmap
-     */
-    public static Bitmap bundleToHardwareBitmap(Bundle bundle) {
-        if (!bundle.containsKey(KEY_BUFFER) || !bundle.containsKey(KEY_COLOR_SPACE)) {
-            throw new IllegalArgumentException("Bundle does not contain a hardware bitmap");
-        }
-
-        HardwareBuffer buffer = bundle.getParcelable(KEY_BUFFER);
-        ParcelableColorSpace colorSpace = bundle.getParcelable(KEY_COLOR_SPACE);
-
-        return Bitmap.wrapHardwareBuffer(Objects.requireNonNull(buffer),
-                colorSpace.getColorSpace());
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
index 98e6bd141b65..924351df3117 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
@@ -51,7 +51,6 @@ import androidx.annotation.NonNull;
 import com.android.internal.logging.UiEventLogger;
 import com.android.internal.util.ScreenshotHelper;
 import com.android.systemui.R;
-import com.android.systemui.shared.recents.utilities.BitmapUtil;
 
 import java.util.function.Consumer;
 
@@ -208,7 +207,7 @@ public class TakeScreenshotService extends Service {
                 if (DEBUG_SERVICE) {
                     Log.d(TAG, "handleMessage: TAKE_SCREENSHOT_PROVIDED_IMAGE");
                 }
-                Bitmap screenshot = BitmapUtil.bundleToHardwareBitmap(
+                Bitmap screenshot = ScreenshotHelper.HardwareBitmapBundler.bundleToHardwareBitmap(
                         screenshotRequest.getBitmapBundle());
                 Rect screenBounds = screenshotRequest.getBoundsInScreen();
                 Insets insets = screenshotRequest.getInsets();
diff --git a/services/core/java/com/android/server/app/GameServiceProviderInstanceFactoryImpl.java b/services/core/java/com/android/server/app/GameServiceProviderInstanceFactoryImpl.java
index 0abab6aafe73..b0a389d8ac68 100644
--- a/services/core/java/com/android/server/app/GameServiceProviderInstanceFactoryImpl.java
+++ b/services/core/java/com/android/server/app/GameServiceProviderInstanceFactoryImpl.java
@@ -29,6 +29,7 @@ import android.service.games.IGameSessionService;
 
 import com.android.internal.infra.ServiceConnector;
 import com.android.internal.os.BackgroundThread;
+import com.android.internal.util.ScreenshotHelper;
 import com.android.server.LocalServices;
 import com.android.server.app.GameServiceConfiguration.GameServiceComponentConfiguration;
 import com.android.server.wm.WindowManagerInternal;
@@ -55,7 +56,8 @@ final class GameServiceProviderInstanceFactoryImpl implements GameServiceProvide
                 (WindowManagerService) ServiceManager.getService(Context.WINDOW_SERVICE),
                 LocalServices.getService(WindowManagerInternal.class),
                 new GameServiceConnector(mContext, configuration),
-                new GameSessionServiceConnector(mContext, configuration));
+                new GameSessionServiceConnector(mContext, configuration),
+                new ScreenshotHelper(mContext));
     }
 
     private static final class GameServiceConnector extends ServiceConnector.Impl<IGameService> {
diff --git a/services/core/java/com/android/server/app/GameServiceProviderInstanceImpl.java b/services/core/java/com/android/server/app/GameServiceProviderInstanceImpl.java
index e4edf4e0f5f9..faf5c3826964 100644
--- a/services/core/java/com/android/server/app/GameServiceProviderInstanceImpl.java
+++ b/services/core/java/com/android/server/app/GameServiceProviderInstanceImpl.java
@@ -29,7 +29,10 @@ import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.graphics.Bitmap;
+import android.graphics.Insets;
 import android.graphics.Rect;
+import android.net.Uri;
+import android.os.Bundle;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.service.games.CreateGameSessionRequest;
@@ -45,12 +48,15 @@ import android.service.games.IGameSessionService;
 import android.util.Slog;
 import android.view.SurfaceControl;
 import android.view.SurfaceControlViewHost.SurfacePackage;
+import android.view.WindowManager;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.infra.AndroidFuture;
 import com.android.internal.infra.ServiceConnector;
 import com.android.internal.infra.ServiceConnector.ServiceLifecycleCallbacks;
+import com.android.internal.os.BackgroundThread;
+import com.android.internal.util.ScreenshotHelper;
 import com.android.server.wm.WindowManagerInternal;
 import com.android.server.wm.WindowManagerInternal.TaskSystemBarsListener;
 import com.android.server.wm.WindowManagerService;
@@ -59,6 +65,7 @@ import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 final class GameServiceProviderInstanceImpl implements GameServiceProviderInstance {
     private static final String TAG = "GameServiceProviderInstance";
@@ -188,6 +195,7 @@ final class GameServiceProviderInstanceImpl implements GameServiceProviderInstan
     private final IActivityTaskManager mActivityTaskManager;
     private final WindowManagerService mWindowManagerService;
     private final WindowManagerInternal mWindowManagerInternal;
+    private final ScreenshotHelper mScreenshotHelper;
     private final ServiceConnector<IGameService> mGameServiceConnector;
     private final ServiceConnector<IGameSessionService> mGameSessionServiceConnector;
 
@@ -207,7 +215,8 @@ final class GameServiceProviderInstanceImpl implements GameServiceProviderInstan
             @NonNull WindowManagerService windowManagerService,
             @NonNull WindowManagerInternal windowManagerInternal,
             @NonNull ServiceConnector<IGameService> gameServiceConnector,
-            @NonNull ServiceConnector<IGameSessionService> gameSessionServiceConnector) {
+            @NonNull ServiceConnector<IGameSessionService> gameSessionServiceConnector,
+            @NonNull ScreenshotHelper screenshotHelper) {
         mUserHandle = userHandle;
         mBackgroundExecutor = backgroundExecutor;
         mContext = context;
@@ -218,6 +227,7 @@ final class GameServiceProviderInstanceImpl implements GameServiceProviderInstan
         mWindowManagerInternal = windowManagerInternal;
         mGameServiceConnector = gameServiceConnector;
         mGameSessionServiceConnector = gameSessionServiceConnector;
+        mScreenshotHelper = screenshotHelper;
     }
 
     @Override
@@ -651,7 +661,26 @@ final class GameServiceProviderInstanceImpl implements GameServiceProviderInstan
                 Slog.w(TAG, "Could not get bitmap for id: " + taskId);
                 callback.complete(GameScreenshotResult.createInternalErrorResult());
             } else {
-                callback.complete(GameScreenshotResult.createSuccessResult(bitmap));
+                final Bundle bundle = ScreenshotHelper.HardwareBitmapBundler.hardwareBitmapToBundle(
+                        bitmap);
+                final RunningTaskInfo runningTaskInfo = getRunningTaskInfoForTask(taskId);
+                if (runningTaskInfo == null) {
+                    Slog.w(TAG, "Could not get running task info for id: " + taskId);
+                    callback.complete(GameScreenshotResult.createInternalErrorResult());
+                }
+                final Rect crop = runningTaskInfo.configuration.windowConfiguration.getBounds();
+                final Consumer<Uri> completionConsumer = (uri) -> {
+                    if (uri == null) {
+                        callback.complete(GameScreenshotResult.createInternalErrorResult());
+                    } else {
+                        callback.complete(GameScreenshotResult.createSuccessResult());
+                    }
+                };
+                mScreenshotHelper.provideScreenshot(bundle, crop, Insets.NONE, taskId,
+                        mUserHandle.getIdentifier(), gameSessionRecord.getComponentName(),
+                        WindowManager.ScreenshotSource.SCREENSHOT_OTHER,
+                        BackgroundThread.getHandler(),
+                        completionConsumer);
             }
         });
     }
diff --git a/services/tests/mockingservicestests/src/android/service/games/GameSessionTest.java b/services/tests/mockingservicestests/src/android/service/games/GameSessionTest.java
index e89c812ba1fb..4fe9cd30e4ff 100644
--- a/services/tests/mockingservicestests/src/android/service/games/GameSessionTest.java
+++ b/services/tests/mockingservicestests/src/android/service/games/GameSessionTest.java
@@ -28,7 +28,6 @@ import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 
-import android.graphics.Bitmap;
 import android.platform.test.annotations.Presubmit;
 import android.service.games.GameSession.ScreenshotCallback;
 import android.testing.AndroidTestingRunner;
@@ -61,7 +60,6 @@ import java.util.concurrent.TimeUnit;
 @TestableLooper.RunWithLooper(setAsMainLooper = true)
 public final class GameSessionTest {
     private static final long WAIT_FOR_CALLBACK_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(1);
-    private static final Bitmap TEST_BITMAP = Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888);
 
     @Mock
     private IGameSessionController mMockGameSessionController;
@@ -101,7 +99,7 @@ public final class GameSessionTest {
                         }
 
                         @Override
-                        public void onSuccess(Bitmap bitmap) {
+                        public void onSuccess() {
                             fail();
                         }
                     });
@@ -131,7 +129,7 @@ public final class GameSessionTest {
                     }
 
                     @Override
-                    public void onSuccess(Bitmap bitmap) {
+                    public void onSuccess() {
                         fail();
                     }
                 });
@@ -160,7 +158,7 @@ public final class GameSessionTest {
                     }
 
                     @Override
-                    public void onSuccess(Bitmap bitmap) {
+                    public void onSuccess() {
                         fail();
                     }
                 });
@@ -170,10 +168,10 @@ public final class GameSessionTest {
     }
 
     @Test
-    public void takeScreenshot_gameManagerSuccess_returnsBitmap() throws Exception {
+    public void takeScreenshot_gameManagerSuccess() throws Exception {
         doAnswer(invocation -> {
             AndroidFuture result = invocation.getArgument(1);
-            result.complete(GameScreenshotResult.createSuccessResult(TEST_BITMAP));
+            result.complete(GameScreenshotResult.createSuccessResult());
             return null;
         }).when(mMockGameSessionController).takeScreenshot(anyInt(), any());
 
@@ -187,8 +185,7 @@ public final class GameSessionTest {
                     }
 
                     @Override
-                    public void onSuccess(Bitmap bitmap) {
-                        assertEquals(TEST_BITMAP, bitmap);
+                    public void onSuccess() {
                         countDownLatch.countDown();
                     }
                 });
diff --git a/services/tests/mockingservicestests/src/com/android/server/app/GameServiceProviderInstanceImplTest.java b/services/tests/mockingservicestests/src/com/android/server/app/GameServiceProviderInstanceImplTest.java
index 32a31d0e57f7..575e3513eae4 100644
--- a/services/tests/mockingservicestests/src/com/android/server/app/GameServiceProviderInstanceImplTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/app/GameServiceProviderInstanceImplTest.java
@@ -46,7 +46,12 @@ import android.content.ContextWrapper;
 import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.graphics.Bitmap;
+import android.graphics.Canvas;
+import android.graphics.Color;
+import android.graphics.Paint;
+import android.graphics.Picture;
 import android.graphics.Rect;
+import android.net.Uri;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.platform.test.annotations.Presubmit;
@@ -71,6 +76,7 @@ import com.android.internal.infra.AndroidFuture;
 import com.android.internal.util.ConcurrentUtils;
 import com.android.internal.util.FunctionalUtils.ThrowingConsumer;
 import com.android.internal.util.Preconditions;
+import com.android.internal.util.ScreenshotHelper;
 import com.android.server.wm.WindowManagerInternal;
 import com.android.server.wm.WindowManagerInternal.TaskSystemBarsListener;
 import com.android.server.wm.WindowManagerService;
@@ -87,6 +93,7 @@ import org.mockito.quality.Strictness;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Objects;
+import java.util.function.Consumer;
 
 
 /**
@@ -114,7 +121,16 @@ public final class GameServiceProviderInstanceImplTest {
             new ComponentName(GAME_B_PACKAGE, "com.package.game.b.MainActivity");
 
 
-    private static final Bitmap TEST_BITMAP = Bitmap.createBitmap(32, 32, Bitmap.Config.ARGB_8888);
+    private static final Bitmap TEST_BITMAP;
+    static {
+        Picture picture = new Picture();
+        Canvas canvas = picture.beginRecording(200, 100);
+        Paint p = new Paint();
+        p.setColor(Color.BLACK);
+        canvas.drawCircle(10, 10, 10, p);
+        picture.endRecording();
+        TEST_BITMAP = Bitmap.createBitmap(picture);
+    }
 
     private MockitoSession mMockingSession;
     private GameServiceProviderInstance mGameServiceProviderInstance;
@@ -126,6 +142,8 @@ public final class GameServiceProviderInstanceImplTest {
     private WindowManagerInternal mMockWindowManagerInternal;
     @Mock
     private IActivityManager mMockActivityManager;
+    @Mock
+    private ScreenshotHelper mMockScreenshotHelper;
     private MockContext mMockContext;
     private FakeGameClassifier mFakeGameClassifier;
     private FakeGameService mFakeGameService;
@@ -192,7 +210,8 @@ public final class GameServiceProviderInstanceImplTest {
                 mMockWindowManagerService,
                 mMockWindowManagerInternal,
                 mFakeGameServiceConnector,
-                mFakeGameSessionServiceConnector);
+                mFakeGameSessionServiceConnector,
+                mMockScreenshotHelper);
     }
 
     @After
@@ -425,6 +444,7 @@ public final class GameServiceProviderInstanceImplTest {
     public void systemBarsTransientShownDueToGesture_hasGameSession_propagatesToGameSession() {
         mGameServiceProviderInstance.start();
         startTask(10, GAME_A_MAIN_ACTIVITY);
+        mockPermissionGranted(Manifest.permission.MANAGE_GAME_ACTIVITY);
         mFakeGameService.requestCreateGameSession(10);
 
         FakeGameSession gameSession10 = new FakeGameSession();
@@ -446,6 +466,7 @@ public final class GameServiceProviderInstanceImplTest {
     public void systemBarsTransientShownButNotGesture_hasGameSession_notPropagatedToGameSession() {
         mGameServiceProviderInstance.start();
         startTask(10, GAME_A_MAIN_ACTIVITY);
+        mockPermissionGranted(Manifest.permission.MANAGE_GAME_ACTIVITY);
         mFakeGameService.requestCreateGameSession(10);
 
         FakeGameSession gameSession10 = new FakeGameSession();
@@ -799,27 +820,32 @@ public final class GameServiceProviderInstanceImplTest {
         SurfaceControl mockOverlaySurfaceControl = Mockito.mock(SurfaceControl.class);
         SurfaceControl[] excludeLayers = new SurfaceControl[1];
         excludeLayers[0] = mockOverlaySurfaceControl;
+        int taskId = 10;
         when(mMockWindowManagerService.captureTaskBitmap(eq(10), any())).thenReturn(TEST_BITMAP);
-
+        doAnswer(invocation -> {
+            Consumer<Uri> consumer = invocation.getArgument(invocation.getArguments().length - 1);
+            consumer.accept(Uri.parse("a/b.png"));
+            return null;
+        }).when(mMockScreenshotHelper).provideScreenshot(
+                any(), any(), any(), anyInt(), anyInt(), any(), anyInt(), any(), any());
         mGameServiceProviderInstance.start();
-        startTask(10, GAME_A_MAIN_ACTIVITY);
+        startTask(taskId, GAME_A_MAIN_ACTIVITY);
         mockPermissionGranted(Manifest.permission.MANAGE_GAME_ACTIVITY);
-        mFakeGameService.requestCreateGameSession(10);
+        mFakeGameService.requestCreateGameSession(taskId);
 
         FakeGameSession gameSession10 = new FakeGameSession();
         SurfacePackage mockOverlaySurfacePackage = Mockito.mock(SurfacePackage.class);
         when(mockOverlaySurfacePackage.getSurfaceControl()).thenReturn(mockOverlaySurfaceControl);
-        mFakeGameSessionService.removePendingFutureForTaskId(10)
+        mFakeGameSessionService.removePendingFutureForTaskId(taskId)
                 .complete(new CreateGameSessionResult(gameSession10, mockOverlaySurfacePackage));
 
         IGameSessionController gameSessionController = getOnlyElement(
                 mFakeGameSessionService.getCapturedCreateInvocations()).mGameSessionController;
         AndroidFuture<GameScreenshotResult> resultFuture = new AndroidFuture<>();
-        gameSessionController.takeScreenshot(10, resultFuture);
+        gameSessionController.takeScreenshot(taskId, resultFuture);
 
         GameScreenshotResult result = resultFuture.get();
         assertEquals(GameScreenshotResult.GAME_SCREENSHOT_SUCCESS, result.getStatus());
-        assertEquals(TEST_BITMAP, result.getBitmap());
     }
 
     @Test