diff --git a/services/surfaceflinger/tests/ScreenCapture_test.cpp b/services/surfaceflinger/tests/ScreenCapture_test.cpp
index 976ee35..013694f 100644
--- a/services/surfaceflinger/tests/ScreenCapture_test.cpp
+++ b/services/surfaceflinger/tests/ScreenCapture_test.cpp
@@ -222,6 +222,14 @@
     mCapture->checkPixel(0, 0, 200, 200, 200);
 }
 
+TEST_F(ScreenCaptureTest, CaptureLayerExcludeThroughDisplayArgs) {
+    mCaptureArgs.excludeHandles = {mFGSurfaceControl->getHandle()};
+    ScreenCapture::captureDisplay(&mCapture, mCaptureArgs);
+    mCapture->expectBGColor(0, 0);
+    // Doesn't capture FG layer which is at 64, 64
+    mCapture->expectBGColor(64, 64);
+}
+
 // Like the last test but verifies that children are also exclude.
 TEST_F(ScreenCaptureTest, CaptureLayerExcludeTree) {
     auto fgHandle = mFGSurfaceControl->getHandle();
diff --git a/services/surfaceflinger/tests/unittests/CompositionTest.cpp b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
index 19a93e1..f596061 100644
--- a/services/surfaceflinger/tests/unittests/CompositionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
@@ -204,7 +204,7 @@
 
     auto traverseLayers = [this](const LayerVector::Visitor& visitor) {
         return mFlinger.traverseLayersInLayerStack(mDisplay->getLayerStack(),
-                                                   CaptureArgs::UNSET_UID, visitor);
+                                                   CaptureArgs::UNSET_UID, {}, visitor);
     };
 
     auto getLayerSnapshots = RenderArea::fromTraverseLayersLambda(traverseLayers);
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index fc9e653..89a190e 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -446,8 +446,10 @@
     }
 
     auto traverseLayersInLayerStack(ui::LayerStack layerStack, int32_t uid,
+                                    std::unordered_set<uint32_t> excludeLayerIds,
                                     const LayerVector::Visitor& visitor) {
-        return mFlinger->SurfaceFlinger::traverseLayersInLayerStack(layerStack, uid, visitor);
+        return mFlinger->SurfaceFlinger::traverseLayersInLayerStack(layerStack, uid,
+                                                                    excludeLayerIds, visitor);
     }
 
     auto getDisplayNativePrimaries(const sp<IBinder>& displayToken,