From f522b54d799f48cbdbc122f69cb8d2530143d43e Mon Sep 17 00:00:00 2001
From: Hawkwood Glazier <jglazier@google.com>
Date: Thu, 9 Mar 2023 16:34:32 +0000
Subject: [PATCH] Refactor ViewScreenshotTestRule slightly to make it a bit
 extensible

Bug: 271331584
Test: atest AODClockOPRTest
Change-Id: I0a2563c5bb8a5e3bbf191eed18b48de8c2214798
---
 .../screenshot/ViewScreenshotTestRule.kt      | 27 ++++++++++++-------
 1 file changed, 17 insertions(+), 10 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 738b37c9eea46..2d47356f2ce16 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
@@ -39,14 +39,14 @@ import platform.test.screenshot.getEmulatedDevicePathConfig
 import platform.test.screenshot.matchers.BitmapMatcher
 
 /** A rule for View screenshot diff unit tests. */
-class ViewScreenshotTestRule(
+open class ViewScreenshotTestRule(
     emulationSpec: DeviceEmulationSpec,
     private val matcher: BitmapMatcher = UnitTestBitmapMatcher,
     assetsPathRelativeToBuildRoot: String
 ) : TestRule {
     private val colorsRule = MaterialYouColorsRule()
     private val deviceEmulationRule = DeviceEmulationRule(emulationSpec)
-    private val screenshotRule =
+    protected val screenshotRule =
         ScreenshotTestRule(
             SystemUIGoldenImagePathManager(
                 getEmulatedDevicePathConfig(emulationSpec),
@@ -64,15 +64,10 @@ class ViewScreenshotTestRule(
         return delegateRule.apply(base, description)
     }
 
-    /**
-     * Compare the content of the view provided by [viewProvider] with the golden image identified
-     * by [goldenIdentifier] in the context of [emulationSpec].
-     */
-    fun screenshotTest(
-        goldenIdentifier: String,
+    protected fun takeScreenshot(
         mode: Mode = Mode.WrapContent,
         viewProvider: (ComponentActivity) -> View,
-    ) {
+    ): Bitmap {
         activityRule.scenario.onActivity { activity ->
             // Make sure that the activity draws full screen and fits the whole display instead of
             // the system bars.
@@ -99,7 +94,19 @@ class ViewScreenshotTestRule(
             contentView = content.getChildAt(0)
         }
 
-        val bitmap = contentView?.toBitmap() ?: error("contentView is null")
+        return contentView?.toBitmap() ?: error("contentView is null")
+    }
+
+    /**
+     * Compare the content of the view provided by [viewProvider] with the golden image identified
+     * by [goldenIdentifier] in the context of [emulationSpec].
+     */
+    fun screenshotTest(
+        goldenIdentifier: String,
+        mode: Mode = Mode.WrapContent,
+        viewProvider: (ComponentActivity) -> View,
+    ) {
+        val bitmap = takeScreenshot(mode, viewProvider)
         screenshotRule.assertBitmapAgainstGolden(
             bitmap,
             goldenIdentifier,