Merge "Adding optional block to run before screenshot" into udc-dev am: a6550a836b

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/23356977

Change-Id: Id62b33726d2ad2e5796200373f64660c930875ce
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewScreenshotTestRule.kt

@@ -78,6 +78,7 @@ open class ViewScreenshotTestRule(
     protected fun takeScreenshot(
         mode: Mode = Mode.WrapContent,
         viewProvider: (ComponentActivity) -> View,
+        beforeScreenshot: (ComponentActivity) -> Unit = {}
     ): Bitmap {
         activityRule.scenario.onActivity { activity ->
             // Make sure that the activity draws full screen and fits the whole display instead of
@@ -103,6 +104,7 @@ open class ViewScreenshotTestRule(
             val content = activity.requireViewById<ViewGroup>(android.R.id.content)
             assertEquals(1, content.childCount)
             contentView = content.getChildAt(0)
+            beforeScreenshot(activity)
         }
 
         return if (isRobolectric) {
@@ -120,9 +122,10 @@ open class ViewScreenshotTestRule(
     fun screenshotTest(
         goldenIdentifier: String,
         mode: Mode = Mode.WrapContent,
-        viewProvider: (ComponentActivity) -> View,
+        beforeScreenshot: (ComponentActivity) -> Unit = {},
+        viewProvider: (ComponentActivity) -> View
     ) {
-        val bitmap = takeScreenshot(mode, viewProvider)
+        val bitmap = takeScreenshot(mode, viewProvider, beforeScreenshot)
         screenshotRule.assertBitmapAgainstGolden(
             bitmap,
             goldenIdentifier,