[ScreenCaptureTest] Remove captureDisplay.

Replace usages of ScreenCapture#captureDisplay with
ScreenCapture#captureLayers. Clean up repetitive
variables and setup that is already performed in
LayerTransactionTest. Move
Surface_test#ScreenshotsOfProtectedBuffersDontSucceed
to where other screenshot tests are and clean up its
logic.

Bug: 293445881
Bug: 303673371
Test: atest ScreenCaptureTest
Test: atest SurfaceTest
Test: atest LayerTransactionTest
Change-Id: Iba7e0f9e3f1b97bf503b3007112850b25867aa86

3 files changed, 76 insertions, 203 deletions

diff --git a/libs/gui/tests/Surface_test.cpp b/libs/gui/tests/Surface_test.cpp
index e7b1232eac..4c18cf74d4 100644
--- a/libs/gui/tests/Surface_test.cpp
+++ b/libs/gui/tests/Surface_test.cpp
@@ -208,21 +208,6 @@ protected:
         ASSERT_EQ(NO_ERROR, surface->disconnect(NATIVE_WINDOW_API_CPU));
     }
 
-    static status_t captureDisplay(DisplayCaptureArgs& captureArgs,
-                                   ScreenCaptureResults& captureResults) {
-        const auto sf = ComposerServiceAIDL::getComposerService();
-        SurfaceComposerClient::Transaction().apply(true);
-
-        const sp<SyncScreenCaptureListener> captureListener = new SyncScreenCaptureListener();
-        binder::Status status = sf->captureDisplay(captureArgs, captureListener);
-        status_t err = gui::aidl_utils::statusTFromBinderStatus(status);
-        if (err != NO_ERROR) {
-            return err;
-        }
-        captureResults = captureListener->waitForResults();
-        return fenceStatus(captureResults.fenceResult);
-    }
-
     sp<Surface> mSurface;
     sp<SurfaceComposerClient> mComposerClient;
     sp<SurfaceControl> mSurfaceControl;
@@ -260,56 +245,6 @@ TEST_F(SurfaceTest, QueuesToWindowComposerIsTrueWhenPurgatorized) {
     EXPECT_EQ(1, result);
 }
 
-// This test probably doesn't belong here.
-TEST_F(SurfaceTest, ScreenshotsOfProtectedBuffersDontSucceed) {
-    sp<ANativeWindow> anw(mSurface);
-
-    // Verify the screenshot works with no protected buffers.
-    const auto ids = SurfaceComposerClient::getPhysicalDisplayIds();
-    ASSERT_FALSE(ids.empty());
-    // display 0 is picked for now, can extend to support all displays if needed
-    const sp<IBinder> display = SurfaceComposerClient::getPhysicalDisplayToken(ids.front());
-    ASSERT_FALSE(display == nullptr);
-
-    DisplayCaptureArgs captureArgs;
-    captureArgs.displayToken = display;
-    captureArgs.width = 64;
-    captureArgs.height = 64;
-
-    ScreenCaptureResults captureResults;
-    ASSERT_EQ(NO_ERROR, captureDisplay(captureArgs, captureResults));
-
-    ASSERT_EQ(NO_ERROR, native_window_api_connect(anw.get(),
-            NATIVE_WINDOW_API_CPU));
-    // Set the PROTECTED usage bit and verify that the screenshot fails.  Note
-    // that we need to dequeue a buffer in order for it to actually get
-    // allocated in SurfaceFlinger.
-    ASSERT_EQ(NO_ERROR, native_window_set_usage(anw.get(),
-            GRALLOC_USAGE_PROTECTED));
-    ASSERT_EQ(NO_ERROR, native_window_set_buffer_count(anw.get(), 3));
-    ANativeWindowBuffer* buf = nullptr;
-
-    status_t err = native_window_dequeue_buffer_and_wait(anw.get(), &buf);
-    if (err) {
-        // we could fail if GRALLOC_USAGE_PROTECTED is not supported.
-        // that's okay as long as this is the reason for the failure.
-        // try again without the GRALLOC_USAGE_PROTECTED bit.
-        ASSERT_EQ(NO_ERROR, native_window_set_usage(anw.get(), 0));
-        ASSERT_EQ(NO_ERROR, native_window_dequeue_buffer_and_wait(anw.get(),
-                &buf));
-        return;
-    }
-    ASSERT_EQ(NO_ERROR, anw->cancelBuffer(anw.get(), buf, -1));
-
-    for (int i = 0; i < 4; i++) {
-        // Loop to make sure SurfaceFlinger has retired a protected buffer.
-        ASSERT_EQ(NO_ERROR, native_window_dequeue_buffer_and_wait(anw.get(),
-                &buf));
-        ASSERT_EQ(NO_ERROR, anw->queueBuffer(anw.get(), buf, -1));
-    }
-    ASSERT_EQ(NO_ERROR, captureDisplay(captureArgs, captureResults));
-}
-
 TEST_F(SurfaceTest, ConcreteTypeIsSurface) {
     sp<ANativeWindow> anw(mSurface);
     int result = -123;
diff --git a/services/surfaceflinger/tests/LayerTransactionTest.h b/services/surfaceflinger/tests/LayerTransactionTest.h
index 2bdb8a452d..076c7e09de 100644
--- a/services/surfaceflinger/tests/LayerTransactionTest.h
+++ b/services/surfaceflinger/tests/LayerTransactionTest.h
@@ -47,7 +47,6 @@ protected:
         ASSERT_NO_FATAL_FAILURE(SetUpDisplay());
 
         sp<gui::ISurfaceComposer> sf(ComposerServiceAIDL::getComposerService());
-        mCaptureArgs.displayToken = mDisplay;
     }
 
     virtual void TearDown() {
@@ -279,8 +278,6 @@ protected:
     const int32_t mLayerZBase = std::numeric_limits<int32_t>::max() - 256;
 
     sp<SurfaceControl> mBlackBgSurface;
-
-    DisplayCaptureArgs mCaptureArgs;
     ScreenCaptureResults mCaptureResults;
 
 private:
diff --git a/services/surfaceflinger/tests/ScreenCapture_test.cpp b/services/surfaceflinger/tests/ScreenCapture_test.cpp
index 96cc333a08..79864e05af 100644
--- a/services/surfaceflinger/tests/ScreenCapture_test.cpp
+++ b/services/surfaceflinger/tests/ScreenCapture_test.cpp
@@ -15,6 +15,8 @@
  */
 
 // TODO(b/129481165): remove the #pragma below and fix conversion issues
+#include <sys/types.h>
+#include <cstdint>
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wconversion"
 
@@ -31,26 +33,22 @@ protected:
         LayerTransactionTest::SetUp();
         ASSERT_EQ(NO_ERROR, mClient->initCheck());
 
-        const auto ids = SurfaceComposerClient::getPhysicalDisplayIds();
-        ASSERT_FALSE(ids.empty());
-        mDisplayToken = SurfaceComposerClient::getPhysicalDisplayToken(ids.front());
-        ASSERT_FALSE(mDisplayToken == nullptr);
-
-        ui::DisplayMode mode;
-        ASSERT_EQ(NO_ERROR, SurfaceComposerClient::getActiveDisplayMode(mDisplayToken, &mode));
-        const ui::Size& resolution = mode.resolution;
-
-        mDisplaySize = resolution;
+        // Root surface
+        mRootSurfaceControl =
+                createLayer(String8("RootTestSurface"), mDisplayWidth, mDisplayHeight, 0);
+        ASSERT_TRUE(mRootSurfaceControl != nullptr);
+        ASSERT_TRUE(mRootSurfaceControl->isValid());
 
         // Background surface
-        mBGSurfaceControl = createLayer(String8("BG Test Surface"), resolution.getWidth(),
-                                        resolution.getHeight(), 0);
+        mBGSurfaceControl = createLayer(String8("BG Test Surface"), mDisplayWidth, mDisplayHeight,
+                                        0, mRootSurfaceControl.get());
         ASSERT_TRUE(mBGSurfaceControl != nullptr);
         ASSERT_TRUE(mBGSurfaceControl->isValid());
         TransactionUtils::fillSurfaceRGBA8(mBGSurfaceControl, 63, 63, 195);
 
         // Foreground surface
-        mFGSurfaceControl = createLayer(String8("FG Test Surface"), 64, 64, 0);
+        mFGSurfaceControl =
+                createLayer(String8("FG Test Surface"), 64, 64, 0, mRootSurfaceControl.get());
 
         ASSERT_TRUE(mFGSurfaceControl != nullptr);
         ASSERT_TRUE(mFGSurfaceControl->isValid());
@@ -58,7 +56,7 @@ protected:
         TransactionUtils::fillSurfaceRGBA8(mFGSurfaceControl, 195, 63, 63);
 
         asTransaction([&](Transaction& t) {
-            t.setDisplayLayerStack(mDisplayToken, ui::DEFAULT_LAYER_STACK);
+            t.setDisplayLayerStack(mDisplay, ui::DEFAULT_LAYER_STACK);
 
             t.setLayer(mBGSurfaceControl, INT32_MAX - 2).show(mBGSurfaceControl);
 
@@ -66,25 +64,22 @@ protected:
                     .setPosition(mFGSurfaceControl, 64, 64)
                     .show(mFGSurfaceControl);
         });
+
+        mCaptureArgs.sourceCrop = mDisplayRect;
+        mCaptureArgs.layerHandle = mRootSurfaceControl->getHandle();
     }
 
     virtual void TearDown() {
         LayerTransactionTest::TearDown();
         mBGSurfaceControl = 0;
         mFGSurfaceControl = 0;
-
-        // Restore display rotation
-        asTransaction([&](Transaction& t) {
-            Rect displayBounds{mDisplaySize};
-            t.setDisplayProjection(mDisplayToken, ui::ROTATION_0, displayBounds, displayBounds);
-        });
     }
 
+    sp<SurfaceControl> mRootSurfaceControl;
     sp<SurfaceControl> mBGSurfaceControl;
     sp<SurfaceControl> mFGSurfaceControl;
     std::unique_ptr<ScreenCapture> mCapture;
-    sp<IBinder> mDisplayToken;
-    ui::Size mDisplaySize;
+    LayerCaptureArgs mCaptureArgs;
 };
 
 TEST_F(ScreenCaptureTest, SetFlagsSecureEUidSystem) {
@@ -92,7 +87,8 @@ TEST_F(ScreenCaptureTest, SetFlagsSecureEUidSystem) {
     ASSERT_NO_FATAL_FAILURE(
             layer = createLayer("test", 32, 32,
                                 ISurfaceComposerClient::eSecure |
-                                        ISurfaceComposerClient::eFXSurfaceBufferQueue));
+                                        ISurfaceComposerClient::eFXSurfaceBufferQueue,
+                                mRootSurfaceControl.get()));
     ASSERT_NO_FATAL_FAILURE(fillBufferQueueLayerColor(layer, Color::RED, 32, 32));
 
     Transaction().show(layer).setLayer(layer, INT32_MAX).apply(true);
@@ -100,14 +96,14 @@ TEST_F(ScreenCaptureTest, SetFlagsSecureEUidSystem) {
     {
         // Ensure the UID is not root because root has all permissions
         UIDFaker f(AID_APP_START);
-        ASSERT_EQ(PERMISSION_DENIED, ScreenCapture::captureDisplay(mCaptureArgs, mCaptureResults));
+        ASSERT_EQ(PERMISSION_DENIED, ScreenCapture::captureLayers(mCaptureArgs, mCaptureResults));
     }
 
     UIDFaker f(AID_SYSTEM);
 
     // By default the system can capture screenshots with secure layers but they
     // will be blacked out
-    ASSERT_EQ(NO_ERROR, ScreenCapture::captureDisplay(mCaptureArgs, mCaptureResults));
+    ASSERT_EQ(NO_ERROR, ScreenCapture::captureLayers(mCaptureArgs, mCaptureResults));
 
     {
         SCOPED_TRACE("as system");
@@ -117,10 +113,8 @@ TEST_F(ScreenCaptureTest, SetFlagsSecureEUidSystem) {
 
     // Here we pass captureSecureLayers = true and since we are AID_SYSTEM we should be able
     // to receive them...we are expected to take care with the results.
-    DisplayCaptureArgs args;
-    args.displayToken = mDisplay;
-    args.captureSecureLayers = true;
-    ASSERT_EQ(NO_ERROR, ScreenCapture::captureDisplay(args, mCaptureResults));
+    mCaptureArgs.captureSecureLayers = true;
+    ASSERT_EQ(NO_ERROR, ScreenCapture::captureLayers(mCaptureArgs, mCaptureResults));
     ASSERT_TRUE(mCaptureResults.capturedSecureLayers);
     ScreenCapture sc(mCaptureResults.buffer, mCaptureResults.capturedHdrLayers);
     sc.expectColor(Rect(0, 0, 32, 32), Color::RED);
@@ -131,7 +125,8 @@ TEST_F(ScreenCaptureTest, CaptureChildSetParentFlagsSecureEUidSystem) {
     ASSERT_NO_FATAL_FAILURE(
             parentLayer = createLayer("parent-test", 32, 32,
                                       ISurfaceComposerClient::eSecure |
-                                              ISurfaceComposerClient::eFXSurfaceBufferQueue));
+                                              ISurfaceComposerClient::eFXSurfaceBufferQueue,
+                                      mRootSurfaceControl.get()));
     ASSERT_NO_FATAL_FAILURE(fillBufferQueueLayerColor(parentLayer, Color::RED, 32, 32));
 
     sp<SurfaceControl> childLayer;
@@ -152,10 +147,8 @@ TEST_F(ScreenCaptureTest, CaptureChildSetParentFlagsSecureEUidSystem) {
 
     // Here we pass captureSecureLayers = true and since we are AID_SYSTEM we should be able
     // to receive them...we are expected to take care with the results.
-    DisplayCaptureArgs args;
-    args.displayToken = mDisplay;
-    args.captureSecureLayers = true;
-    ASSERT_EQ(NO_ERROR, ScreenCapture::captureDisplay(args, mCaptureResults));
+    mCaptureArgs.captureSecureLayers = true;
+    ASSERT_EQ(NO_ERROR, ScreenCapture::captureLayers(mCaptureArgs, mCaptureResults));
     ASSERT_TRUE(mCaptureResults.capturedSecureLayers);
     ScreenCapture sc(mCaptureResults.buffer, mCaptureResults.capturedHdrLayers);
     sc.expectColor(Rect(0, 0, 10, 10), Color::BLUE);
@@ -232,7 +225,7 @@ TEST_F(ScreenCaptureTest, CaptureLayerExclude) {
 
 TEST_F(ScreenCaptureTest, CaptureLayerExcludeThroughDisplayArgs) {
     mCaptureArgs.excludeHandles = {mFGSurfaceControl->getHandle()};
-    ScreenCapture::captureDisplay(&mCapture, mCaptureArgs);
+    ScreenCapture::captureLayers(&mCapture, mCaptureArgs);
     mCapture->expectBGColor(0, 0);
     // Doesn't capture FG layer which is at 64, 64
     mCapture->expectBGColor(64, 64);
@@ -605,60 +598,55 @@ TEST_F(ScreenCaptureTest, CaptureSecureLayer) {
     mCapture->expectColor(Rect(30, 30, 60, 60), Color::RED);
 }
 
-TEST_F(ScreenCaptureTest, CaptureDisplayWithUid) {
-    uid_t fakeUid = 12345;
+TEST_F(ScreenCaptureTest, ScreenshotProtectedBuffer) {
+    const uint32_t bufferWidth = 60;
+    const uint32_t bufferHeight = 60;
 
-    DisplayCaptureArgs captureArgs;
-    captureArgs.displayToken = mDisplay;
+    sp<SurfaceControl> layer =
+            createLayer(String8("Colored surface"), bufferWidth, bufferHeight,
+                        ISurfaceComposerClient::eFXSurfaceBufferState, mRootSurfaceControl.get());
 
-    sp<SurfaceControl> layer;
-    ASSERT_NO_FATAL_FAILURE(layer = createLayer("test layer", 32, 32,
-                                                ISurfaceComposerClient::eFXSurfaceBufferQueue,
-                                                mBGSurfaceControl.get()));
-    ASSERT_NO_FATAL_FAILURE(fillBufferQueueLayerColor(layer, Color::RED, 32, 32));
+    Transaction().show(layer).setLayer(layer, INT32_MAX).apply(true);
 
-    Transaction().show(layer).setLayer(layer, INT32_MAX).apply();
+    sp<Surface> surface = layer->getSurface();
+    ASSERT_TRUE(surface != nullptr);
+    sp<ANativeWindow> anw(surface);
 
-    // Make sure red layer with the background layer is screenshot.
-    ScreenCapture::captureDisplay(&mCapture, captureArgs);
-    mCapture->expectColor(Rect(0, 0, 32, 32), Color::RED);
-    mCapture->expectBorder(Rect(0, 0, 32, 32), {63, 63, 195, 255});
+    ASSERT_EQ(NO_ERROR, native_window_api_connect(anw.get(), NATIVE_WINDOW_API_CPU));
+    ASSERT_EQ(NO_ERROR, native_window_set_usage(anw.get(), GRALLOC_USAGE_PROTECTED));
 
-    // From non system uid, can't request screenshot without a specified uid.
-    UIDFaker f(fakeUid);
-    ASSERT_EQ(PERMISSION_DENIED, ScreenCapture::captureDisplay(captureArgs, mCaptureResults));
+    int fenceFd;
+    ANativeWindowBuffer* buf = nullptr;
 
-    // Make screenshot request with current uid set. No layers were created with the current
-    // uid so screenshot will be black.
-    captureArgs.uid = fakeUid;
-    ScreenCapture::captureDisplay(&mCapture, captureArgs);
-    mCapture->expectColor(Rect(0, 0, 32, 32), Color::BLACK);
-    mCapture->expectBorder(Rect(0, 0, 32, 32), Color::BLACK);
+    // End test if device does not support USAGE_PROTECTED
+    // b/309965549 This check does not exit the test when running on AVDs
+    status_t err = anw->dequeueBuffer(anw.get(), &buf, &fenceFd);
+    if (err) {
+        return;
+    }
+    anw->queueBuffer(anw.get(), buf, fenceFd);
 
-    sp<SurfaceControl> layerWithFakeUid;
-    // Create a new layer with the current uid
-    ASSERT_NO_FATAL_FAILURE(layerWithFakeUid =
-                                    createLayer("new test layer", 32, 32,
-                                                ISurfaceComposerClient::eFXSurfaceBufferQueue,
-                                                mBGSurfaceControl.get()));
-    ASSERT_NO_FATAL_FAILURE(fillBufferQueueLayerColor(layerWithFakeUid, Color::GREEN, 32, 32));
-    Transaction()
-            .show(layerWithFakeUid)
-            .setLayer(layerWithFakeUid, INT32_MAX)
-            .setPosition(layerWithFakeUid, 128, 128)
-            .apply();
+    // USAGE_PROTECTED buffer is read as a black screen
+    ScreenCaptureResults captureResults;
+    ASSERT_EQ(NO_ERROR, ScreenCapture::captureLayers(mCaptureArgs, captureResults));
 
-    // Screenshot from the fakeUid caller with the uid requested allows the layer
-    // with that uid to be screenshotted. Everything else is black
-    ScreenCapture::captureDisplay(&mCapture, captureArgs);
-    mCapture->expectColor(Rect(128, 128, 160, 160), Color::GREEN);
-    mCapture->expectBorder(Rect(128, 128, 160, 160), Color::BLACK);
+    ScreenCapture sc(captureResults.buffer, captureResults.capturedHdrLayers);
+    sc.expectColor(Rect(0, 0, bufferWidth, bufferHeight), Color::BLACK);
+
+    // Reading color data will expectedly result in crash, only check usage bit
+    // b/309965549 Checking that the usage bit is protected does not work for
+    // devices that do not support usage protected.
+    mCaptureArgs.allowProtected = true;
+    ASSERT_EQ(NO_ERROR, ScreenCapture::captureLayers(mCaptureArgs, captureResults));
+    // ASSERT_EQ(GRALLOC_USAGE_PROTECTED, GRALLOC_USAGE_PROTECTED &
+    // captureResults.buffer->getUsage());
 }
 
-TEST_F(ScreenCaptureTest, CaptureDisplayPrimaryDisplayOnly) {
+TEST_F(ScreenCaptureTest, CaptureLayer) {
     sp<SurfaceControl> layer;
-    ASSERT_NO_FATAL_FAILURE(
-            layer = createLayer("test layer", 0, 0, ISurfaceComposerClient::eFXSurfaceEffect));
+    ASSERT_NO_FATAL_FAILURE(layer = createLayer("test layer", 0, 0,
+                                                ISurfaceComposerClient::eFXSurfaceEffect,
+                                                mRootSurfaceControl.get()));
 
     const Color layerColor = Color::RED;
     const Rect bounds = Rect(10, 10, 40, 40);
@@ -666,17 +654,13 @@ TEST_F(ScreenCaptureTest, CaptureDisplayPrimaryDisplayOnly) {
     Transaction()
             .show(layer)
             .hide(mFGSurfaceControl)
-            .setLayerStack(layer, ui::DEFAULT_LAYER_STACK)
             .setLayer(layer, INT32_MAX)
             .setColor(layer, {layerColor.r / 255, layerColor.g / 255, layerColor.b / 255})
             .setCrop(layer, bounds)
             .apply();
 
-    DisplayCaptureArgs captureArgs;
-    captureArgs.displayToken = mDisplay;
-
     {
-        ScreenCapture::captureDisplay(&mCapture, captureArgs);
+        ScreenCapture::captureLayers(&mCapture, mCaptureArgs);
         mCapture->expectColor(bounds, layerColor);
         mCapture->expectBorder(bounds, {63, 63, 195, 255});
     }
@@ -689,17 +673,18 @@ TEST_F(ScreenCaptureTest, CaptureDisplayPrimaryDisplayOnly) {
     {
         // Can't screenshot test layer since it now has flag
         // eLayerSkipScreenshot
-        ScreenCapture::captureDisplay(&mCapture, captureArgs);
+        ScreenCapture::captureLayers(&mCapture, mCaptureArgs);
         mCapture->expectColor(bounds, {63, 63, 195, 255});
         mCapture->expectBorder(bounds, {63, 63, 195, 255});
     }
 }
 
-TEST_F(ScreenCaptureTest, CaptureDisplayChildPrimaryDisplayOnly) {
+TEST_F(ScreenCaptureTest, CaptureLayerChild) {
     sp<SurfaceControl> layer;
     sp<SurfaceControl> childLayer;
-    ASSERT_NO_FATAL_FAILURE(
-            layer = createLayer("test layer", 0, 0, ISurfaceComposerClient::eFXSurfaceEffect));
+    ASSERT_NO_FATAL_FAILURE(layer = createLayer("test layer", 0, 0,
+                                                ISurfaceComposerClient::eFXSurfaceEffect,
+                                                mRootSurfaceControl.get()));
     ASSERT_NO_FATAL_FAILURE(childLayer = createLayer("test layer", 0, 0,
                                                      ISurfaceComposerClient::eFXSurfaceEffect,
                                                      layer.get()));
@@ -713,7 +698,6 @@ TEST_F(ScreenCaptureTest, CaptureDisplayChildPrimaryDisplayOnly) {
             .show(layer)
             .show(childLayer)
             .hide(mFGSurfaceControl)
-            .setLayerStack(layer, ui::DEFAULT_LAYER_STACK)
             .setLayer(layer, INT32_MAX)
             .setColor(layer, {layerColor.r / 255, layerColor.g / 255, layerColor.b / 255})
             .setColor(childLayer, {childColor.r / 255, childColor.g / 255, childColor.b / 255})
@@ -721,11 +705,8 @@ TEST_F(ScreenCaptureTest, CaptureDisplayChildPrimaryDisplayOnly) {
             .setCrop(childLayer, childBounds)
             .apply();
 
-    DisplayCaptureArgs captureArgs;
-    captureArgs.displayToken = mDisplay;
-
     {
-        ScreenCapture::captureDisplay(&mCapture, captureArgs);
+        ScreenCapture::captureLayers(&mCapture, mCaptureArgs);
         mCapture->expectColor(childBounds, childColor);
         mCapture->expectBorder(childBounds, layerColor);
         mCapture->expectBorder(bounds, {63, 63, 195, 255});
@@ -739,7 +720,7 @@ TEST_F(ScreenCaptureTest, CaptureDisplayChildPrimaryDisplayOnly) {
     {
         // Can't screenshot child layer since the parent has the flag
         // eLayerSkipScreenshot
-        ScreenCapture::captureDisplay(&mCapture, captureArgs);
+        ScreenCapture::captureLayers(&mCapture, mCaptureArgs);
         mCapture->expectColor(childBounds, {63, 63, 195, 255});
         mCapture->expectBorder(childBounds, {63, 63, 195, 255});
         mCapture->expectBorder(bounds, {63, 63, 195, 255});
@@ -860,14 +841,10 @@ TEST_F(ScreenCaptureTest, CaptureOffscreen) {
 
     Transaction().show(layer).hide(mFGSurfaceControl).reparent(layer, nullptr).apply();
 
-    DisplayCaptureArgs displayCaptureArgs;
-    displayCaptureArgs.displayToken = mDisplay;
-
     {
         // Validate that the red layer is not on screen
-        ScreenCapture::captureDisplay(&mCapture, displayCaptureArgs);
-        mCapture->expectColor(Rect(0, 0, mDisplaySize.width, mDisplaySize.height),
-                              {63, 63, 195, 255});
+        ScreenCapture::captureLayers(&mCapture, mCaptureArgs);
+        mCapture->expectColor(Rect(0, 0, mDisplayWidth, mDisplayHeight), {63, 63, 195, 255});
     }
 
     LayerCaptureArgs captureArgs;
@@ -878,42 +855,6 @@ TEST_F(ScreenCaptureTest, CaptureOffscreen) {
     mCapture->expectColor(Rect(0, 0, 32, 32), Color::RED);
 }
 
-TEST_F(ScreenCaptureTest, CaptureDisplayWith90DegRotation) {
-    asTransaction([&](Transaction& t) {
-        Rect newDisplayBounds{mDisplaySize.height, mDisplaySize.width};
-        t.setDisplayProjection(mDisplayToken, ui::ROTATION_90, newDisplayBounds, newDisplayBounds);
-    });
-
-    DisplayCaptureArgs displayCaptureArgs;
-    displayCaptureArgs.displayToken = mDisplayToken;
-    displayCaptureArgs.width = mDisplaySize.width;
-    displayCaptureArgs.height = mDisplaySize.height;
-    displayCaptureArgs.useIdentityTransform = true;
-    ScreenCapture::captureDisplay(&mCapture, displayCaptureArgs);
-
-    mCapture->expectBGColor(0, 0);
-    mCapture->expectFGColor(mDisplaySize.width - 65, 65);
-}
-
-TEST_F(ScreenCaptureTest, CaptureDisplayWith270DegRotation) {
-    asTransaction([&](Transaction& t) {
-        Rect newDisplayBounds{mDisplaySize.height, mDisplaySize.width};
-        t.setDisplayProjection(mDisplayToken, ui::ROTATION_270, newDisplayBounds, newDisplayBounds);
-    });
-
-    DisplayCaptureArgs displayCaptureArgs;
-    displayCaptureArgs.displayToken = mDisplayToken;
-    displayCaptureArgs.width = mDisplaySize.width;
-    displayCaptureArgs.height = mDisplaySize.height;
-    displayCaptureArgs.useIdentityTransform = true;
-    ScreenCapture::captureDisplay(&mCapture, displayCaptureArgs);
-
-    std::this_thread::sleep_for(std::chrono::seconds{5});
-
-    mCapture->expectBGColor(mDisplayWidth - 1, mDisplaySize.height - 1);
-    mCapture->expectFGColor(65, mDisplaySize.height - 65);
-}
-
 TEST_F(ScreenCaptureTest, CaptureNonHdrLayer) {
     sp<SurfaceControl> layer;
     ASSERT_NO_FATAL_FAILURE(layer = createLayer("test layer", 32, 32,