diff options
| -rw-r--r-- | packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewScreenshotTestRule.kt | 7 |
1 files changed, 5 insertions, 2 deletions
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewScreenshotTestRule.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewScreenshotTestRule.kt index f96d1e3d7fef..ba73bd7c2d16 100644 --- a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewScreenshotTestRule.kt +++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewScreenshotTestRule.kt @@ -69,6 +69,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 @@ -94,6 +95,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) { @@ -111,9 +113,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, |