diff --git a/packages/SystemUI/screenshot/Android.bp b/packages/SystemUI/screenshot/Android.bp
index a79fd9040d..601e92f 100644
--- a/packages/SystemUI/screenshot/Android.bp
+++ b/packages/SystemUI/screenshot/Android.bp
@@ -26,11 +26,7 @@
     manifest: "AndroidManifest.xml",
 
     srcs: [
-        // All files in this library should be in Kotlin besides some exceptions.
         "src/**/*.kt",
-
-        // This file was forked from google3, so exceptionally it can be in Java.
-        "src/com/android/systemui/testing/screenshot/DynamicColorsTestUtils.java",
     ],
 
     resource_dirs: [
diff --git a/packages/SystemUI/screenshot/AndroidManifest.xml b/packages/SystemUI/screenshot/AndroidManifest.xml
index 3b703be..a405836 100644
--- a/packages/SystemUI/screenshot/AndroidManifest.xml
+++ b/packages/SystemUI/screenshot/AndroidManifest.xml
@@ -23,6 +23,4 @@
             android:exported="true"
             android:theme="@style/Theme.SystemUI.Screenshot" />
     </application>
-
-    <uses-permission android:name="android.permission.WRITE_SECURE_SETTINGS" />
 </manifest>
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/Bitmap.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/Bitmap.kt
new file mode 100644
index 0000000..3d26cda
--- /dev/null
+++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/Bitmap.kt
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.testing.screenshot
+
+import android.graphics.Bitmap
+import android.graphics.Canvas
+import android.os.Build
+import android.view.View
+import platform.test.screenshot.matchers.MSSIMMatcher
+import platform.test.screenshot.matchers.PixelPerfectMatcher
+
+/** Draw this [View] into a [Bitmap]. */
+fun View.drawIntoBitmap(): Bitmap {
+    val bitmap =
+        Bitmap.createBitmap(
+            measuredWidth,
+            measuredHeight,
+            Bitmap.Config.ARGB_8888,
+        )
+    val canvas = Canvas(bitmap)
+    draw(canvas)
+    return bitmap
+}
+
+/**
+ * The [BitmapMatcher][platform.test.screenshot.matchers.BitmapMatcher] that should be used for
+ * screenshot *unit* tests.
+ */
+val UnitTestBitmapMatcher =
+    if (Build.CPU_ABI == "x86_64") {
+        // Different CPU architectures can sometimes end up rendering differently, so we can't do
+        // pixel-perfect matching on different architectures using the same golden. Given that our
+        // presubmits are run on cf_x86_64_phone, our goldens should be perfectly matched on the
+        // x86_64 architecture and use the Structural Similarity Index on others.
+        // TODO(b/237511747): Run our screenshot presubmit tests on arm64 instead so that we can
+        // do pixel perfect matching both at presubmit time and at development time with actual
+        // devices.
+        PixelPerfectMatcher()
+    } else {
+        MSSIMMatcher()
+    }
+
+/**
+ * The [BitmapMatcher][platform.test.screenshot.matchers.BitmapMatcher] that should be used for
+ * screenshot *unit* tests.
+ *
+ * We use the Structural Similarity Index for integration tests because they usually contain
+ * additional information and noise that shouldn't break the test.
+ */
+val IntegrationTestBitmapMatcher = MSSIMMatcher()
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/DynamicColorsTestUtils.java b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/DynamicColorsTestUtils.java
deleted file mode 100644
index 96ec4c5..0000000
--- a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/DynamicColorsTestUtils.java
+++ /dev/null
@@ -1,193 +0,0 @@
-/*
- * Copyright (C) 2022 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.testing.screenshot;
-
-import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
-
-import android.app.UiAutomation;
-import android.content.Context;
-import android.provider.Settings;
-import android.util.Log;
-
-import androidx.annotation.ColorInt;
-import androidx.annotation.ColorRes;
-import androidx.annotation.NonNull;
-import androidx.annotation.RequiresApi;
-import androidx.core.content.ContextCompat;
-import androidx.test.espresso.Espresso;
-import androidx.test.espresso.IdlingRegistry;
-import androidx.test.espresso.IdlingResource;
-
-import org.json.JSONObject;
-import org.junit.function.ThrowingRunnable;
-
-import java.util.HashMap;
-import java.util.Map;
-
-/*
- * Note: This file was forked from
- * google3/third_party/java_src/android_libs/material_components/screenshot_tests/java/android/
- * support/design/scuba/color/DynamicColorsTestUtils.java.
- */
-
-/** Utility that helps change the dynamic system colors for testing. */
-@RequiresApi(32)
-public class DynamicColorsTestUtils {
-
-    private static final String TAG = DynamicColorsTestUtils.class.getSimpleName();
-
-    private static final String THEME_CUSTOMIZATION_KEY = "theme_customization_overlay_packages";
-    private static final String THEME_CUSTOMIZATION_SYSTEM_PALETTE_KEY =
-            "android.theme.customization.system_palette";
-
-    private static final int ORANGE_SYSTEM_SEED_COLOR = 0xA66800;
-    private static final int ORANGE_EXPECTED_SYSTEM_ACCENT1_600_COLOR = -8235756;
-
-    private DynamicColorsTestUtils() {
-    }
-
-    /**
-     * Update system dynamic colors (e.g., android.R.color.system_accent1_600) based on an orange
-     * seed color, and then wait for the change to propagate to the app by comparing
-     * android.R.color.system_accent1_600 to the expected orange value.
-     */
-    public static void updateSystemColorsToOrange() {
-        updateSystemColors(ORANGE_SYSTEM_SEED_COLOR, ORANGE_EXPECTED_SYSTEM_ACCENT1_600_COLOR);
-    }
-
-    /**
-     * Update system dynamic colors (e.g., android.R.color.system_accent1_600) based on the provided
-     * {@code seedColor}, and then wait for the change to propagate to the app by comparing
-     * android.R.color.system_accent1_600 to {@code expectedSystemAccent1600}.
-     */
-    public static void updateSystemColors(
-            @ColorInt int seedColor, @ColorInt int expectedSystemAccent1600) {
-        Context context = getInstrumentation().getTargetContext();
-
-        int actualSystemAccent1600 =
-                ContextCompat.getColor(context, android.R.color.system_accent1_600);
-
-        if (expectedSystemAccent1600 == actualSystemAccent1600) {
-            String expectedColorString = Integer.toHexString(expectedSystemAccent1600);
-            Log.d(
-                    TAG,
-                    "Skipped updating system colors since system_accent1_600 is already equal to "
-                            + "expected: "
-                            + expectedColorString);
-            return;
-        }
-
-        updateSystemColors(seedColor);
-    }
-
-    /**
-     * Update system dynamic colors (e.g., android.R.color.system_accent1_600) based on the provided
-     * {@code seedColor}, and then wait for the change to propagate to the app by checking
-     * android.R.color.system_accent1_600 for any change.
-     */
-    public static void updateSystemColors(@ColorInt int seedColor) {
-        Context context = getInstrumentation().getTargetContext();
-
-        // Initialize system color idling resource with original system_accent1_600 value.
-        ColorChangeIdlingResource systemColorIdlingResource =
-                new ColorChangeIdlingResource(context, android.R.color.system_accent1_600);
-
-        // Update system theme color setting to trigger fabricated resource overlay.
-        runWithShellPermissionIdentity(
-                () ->
-                        Settings.Secure.putString(
-                                context.getContentResolver(),
-                                THEME_CUSTOMIZATION_KEY,
-                                buildThemeCustomizationString(seedColor)));
-
-        // Wait for system color update to propagate to app.
-        IdlingRegistry idlingRegistry = IdlingRegistry.getInstance();
-        idlingRegistry.register(systemColorIdlingResource);
-        Espresso.onIdle();
-        idlingRegistry.unregister(systemColorIdlingResource);
-
-        Log.d(TAG,
-                Settings.Secure.getString(context.getContentResolver(), THEME_CUSTOMIZATION_KEY));
-    }
-
-    private static String buildThemeCustomizationString(@ColorInt int seedColor) {
-        String seedColorHex = Integer.toHexString(seedColor);
-        Map<String, String> themeCustomizationMap = new HashMap<>();
-        themeCustomizationMap.put(THEME_CUSTOMIZATION_SYSTEM_PALETTE_KEY, seedColorHex);
-        return new JSONObject(themeCustomizationMap).toString();
-    }
-
-    private static void runWithShellPermissionIdentity(@NonNull ThrowingRunnable runnable) {
-        UiAutomation uiAutomation = getInstrumentation().getUiAutomation();
-        uiAutomation.adoptShellPermissionIdentity();
-        try {
-            runnable.run();
-        } catch (Throwable e) {
-            throw new RuntimeException(e);
-        } finally {
-            uiAutomation.dropShellPermissionIdentity();
-        }
-    }
-
-    private static class ColorChangeIdlingResource implements IdlingResource {
-
-        private final Context mContext;
-        private final int mColorResId;
-        private final int mInitialColorInt;
-
-        private ResourceCallback mResourceCallback;
-        private boolean mIdleNow;
-
-        ColorChangeIdlingResource(Context context, @ColorRes int colorResId) {
-            this.mContext = context;
-            this.mColorResId = colorResId;
-            this.mInitialColorInt = ContextCompat.getColor(context, colorResId);
-        }
-
-        @Override
-        public String getName() {
-            return ColorChangeIdlingResource.class.getName();
-        }
-
-        @Override
-        public boolean isIdleNow() {
-            if (mIdleNow) {
-                return true;
-            }
-
-            int currentColorInt = ContextCompat.getColor(mContext, mColorResId);
-
-            String initialColorString = Integer.toHexString(mInitialColorInt);
-            String currentColorString = Integer.toHexString(currentColorInt);
-            Log.d(TAG, String.format("Initial=%s, Current=%s", initialColorString,
-                    currentColorString));
-
-            mIdleNow = currentColorInt != mInitialColorInt;
-            Log.d(TAG, String.format("idleNow=%b", mIdleNow));
-
-            if (mIdleNow) {
-                mResourceCallback.onTransitionToIdle();
-            }
-            return mIdleNow;
-        }
-
-        @Override
-        public void registerIdleTransitionCallback(ResourceCallback resourceCallback) {
-            this.mResourceCallback = resourceCallback;
-        }
-    }
-}
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ScreenshotTestRule.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ScreenshotTestRule.kt
deleted file mode 100644
index 564901c..0000000
--- a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ScreenshotTestRule.kt
+++ /dev/null
@@ -1,223 +0,0 @@
-/*
- * Copyright (C) 2022 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.testing.screenshot
-
-import android.app.UiModeManager
-import android.content.Context
-import android.graphics.Bitmap
-import android.graphics.Canvas
-import android.os.Build
-import android.os.UserHandle
-import android.view.Display
-import android.view.View
-import android.view.WindowManagerGlobal
-import androidx.test.platform.app.InstrumentationRegistry
-import org.junit.rules.TestRule
-import org.junit.runner.Description
-import org.junit.runners.model.Statement
-import platform.test.screenshot.GoldenImagePathManager
-import platform.test.screenshot.PathConfig
-import platform.test.screenshot.PathElementNoContext
-import platform.test.screenshot.ScreenshotTestRule
-import platform.test.screenshot.matchers.MSSIMMatcher
-import platform.test.screenshot.matchers.PixelPerfectMatcher
-
-/**
- * A base rule for screenshot diff tests.
- *
- * This rules takes care of setting up the activity according to [testSpec] by:
- * - emulating the display size and density.
- * - setting the dark/light mode.
- * - setting the system (Material You) colors to a fixed value.
- *
- * @see ComposeScreenshotTestRule
- * @see ViewScreenshotTestRule
- */
-class ScreenshotTestRule(private val testSpec: ScreenshotTestSpec) : TestRule {
-    private var currentDisplay: DisplaySpec? = null
-    private var currentGoldenIdentifier: String? = null
-
-    private val pathConfig =
-        PathConfig(
-            PathElementNoContext("model", isDir = true) {
-                currentDisplay?.name ?: error("currentDisplay is null")
-            },
-        )
-    private val matcher = if (shouldUsePerfectMatching()) {
-        PixelPerfectMatcher()
-    } else {
-        MSSIMMatcher()
-    }
-
-    private val screenshotRule =
-        ScreenshotTestRule(
-            SystemUIGoldenImagePathManager(
-                pathConfig,
-                currentGoldenIdentifier = {
-                    currentGoldenIdentifier ?: error("currentGoldenIdentifier is null")
-                },
-            )
-        )
-
-    private fun shouldUsePerfectMatching(): Boolean {
-        // Different CPU architectures can sometimes end up rendering differently, so we can't do
-        // pixel-perfect matching on different architectures using the same golden. Given that our
-        // presubmits are run on cf_x86_64_phone, our goldens should be perfectly matched on the
-        // x86_64 architecture and use the Structural Similarity Index on others.
-        // TODO(b/237511747): Run our screenshot presubmit tests on arm64 instead so that we can
-        // do pixel perfect matching both at presubmit time and at development time with actual
-        // devices.
-        return Build.CPU_ABI == "x86_64"
-    }
-
-    override fun apply(base: Statement, description: Description): Statement {
-        // The statement which call beforeTest() before running the test and afterTest() afterwards.
-        val statement =
-            object : Statement() {
-                override fun evaluate() {
-                    try {
-                        beforeTest()
-                        base.evaluate()
-                    } finally {
-                        afterTest()
-                    }
-                }
-            }
-
-        return screenshotRule.apply(statement, description)
-    }
-
-    private fun beforeTest() {
-        // Update the system colors to a fixed color, so that tests don't depend on the host device
-        // extracted colors. Note that we don't restore the default device colors at the end of the
-        // test because changing the colors (and waiting for them to be applied) is costly and makes
-        // the screenshot tests noticeably slower.
-        DynamicColorsTestUtils.updateSystemColorsToOrange()
-
-        // Emulate the display size and density.
-        val display = testSpec.display
-        val density = display.densityDpi
-        val wm = WindowManagerGlobal.getWindowManagerService()
-        val (width, height) = getEmulatedDisplaySize()
-        wm.setForcedDisplayDensityForUser(Display.DEFAULT_DISPLAY, density, UserHandle.myUserId())
-        wm.setForcedDisplaySize(Display.DEFAULT_DISPLAY, width, height)
-
-        // Force the dark/light theme.
-        val uiModeManager =
-            InstrumentationRegistry.getInstrumentation()
-                .targetContext
-                .getSystemService(Context.UI_MODE_SERVICE) as UiModeManager
-        uiModeManager.setApplicationNightMode(
-            if (testSpec.isDarkTheme) {
-                UiModeManager.MODE_NIGHT_YES
-            } else {
-                UiModeManager.MODE_NIGHT_NO
-            }
-        )
-    }
-
-    private fun afterTest() {
-        // Reset the density and display size.
-        val wm = WindowManagerGlobal.getWindowManagerService()
-        wm.clearForcedDisplayDensityForUser(Display.DEFAULT_DISPLAY, UserHandle.myUserId())
-        wm.clearForcedDisplaySize(Display.DEFAULT_DISPLAY)
-
-        // Reset the dark/light theme.
-        val uiModeManager =
-            InstrumentationRegistry.getInstrumentation()
-                .targetContext
-                .getSystemService(Context.UI_MODE_SERVICE) as UiModeManager
-        uiModeManager.setApplicationNightMode(UiModeManager.MODE_NIGHT_AUTO)
-    }
-
-    /**
-     * Compare the content of [view] with the golden image identified by [goldenIdentifier] in the
-     * context of [testSpec].
-     */
-    fun screenshotTest(goldenIdentifier: String, view: View) {
-        val bitmap = drawIntoBitmap(view)
-
-        // Compare bitmap against golden asset.
-        val isDarkTheme = testSpec.isDarkTheme
-        val isLandscape = testSpec.isLandscape
-        val identifierWithSpec = buildString {
-            append(goldenIdentifier)
-            if (isDarkTheme) append("_dark")
-            if (isLandscape) append("_landscape")
-        }
-
-        // TODO(b/230832101): Provide a way to pass a PathConfig and override the file name on
-        // device to assertBitmapAgainstGolden instead?
-        currentDisplay = testSpec.display
-        currentGoldenIdentifier = goldenIdentifier
-        screenshotRule.assertBitmapAgainstGolden(bitmap, identifierWithSpec, matcher)
-        currentDisplay = null
-        currentGoldenIdentifier = goldenIdentifier
-    }
-
-    /** Draw [view] into a [Bitmap]. */
-    private fun drawIntoBitmap(view: View): Bitmap {
-        val bitmap =
-            Bitmap.createBitmap(
-                view.measuredWidth,
-                view.measuredHeight,
-                Bitmap.Config.ARGB_8888,
-            )
-        val canvas = Canvas(bitmap)
-        view.draw(canvas)
-        return bitmap
-    }
-
-    /** Get the emulated display size for [testSpec]. */
-    private fun getEmulatedDisplaySize(): Pair<Int, Int> {
-        val display = testSpec.display
-        val isPortraitNaturalPosition = display.width < display.height
-        return if (testSpec.isLandscape) {
-            if (isPortraitNaturalPosition) {
-                display.height to display.width
-            } else {
-                display.width to display.height
-            }
-        } else {
-            if (isPortraitNaturalPosition) {
-                display.width to display.height
-            } else {
-                display.height to display.width
-            }
-        }
-    }
-}
-
-private class SystemUIGoldenImagePathManager(
-    pathConfig: PathConfig,
-    private val currentGoldenIdentifier: () -> String,
-) :
-    GoldenImagePathManager(
-        appContext = InstrumentationRegistry.getInstrumentation().context,
-        deviceLocalPath =
-            InstrumentationRegistry.getInstrumentation()
-                .targetContext
-                .filesDir
-                .absolutePath
-                .toString() + "/sysui_screenshots",
-        pathConfig = pathConfig,
-    ) {
-    // This string is appended to all actual/expected screenshots on the device. We append the
-    // golden identifier so that our pull_golden.py scripts can map a screenshot on device to its
-    // asset (and automatically update it, if necessary).
-    override fun toString() = currentGoldenIdentifier()
-}
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ScreenshotTestSpec.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ScreenshotTestSpec.kt
deleted file mode 100644
index 7fc6245..0000000
--- a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ScreenshotTestSpec.kt
+++ /dev/null
@@ -1,78 +0,0 @@
-/*
- * Copyright (C) 2022 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.testing.screenshot
-
-/** The specification of a device display to be used in a screenshot test. */
-data class DisplaySpec(
-    val name: String,
-    val width: Int,
-    val height: Int,
-    val densityDpi: Int,
-)
-
-/** The specification of a screenshot diff test. */
-class ScreenshotTestSpec(
-    val display: DisplaySpec,
-    val isDarkTheme: Boolean = false,
-    val isLandscape: Boolean = false,
-) {
-    companion object {
-        /**
-         * Return a list of [ScreenshotTestSpec] for each of the [displays].
-         *
-         * If [isDarkTheme] is null, this will create a spec for both light and dark themes, for
-         * each of the orientation.
-         *
-         * If [isLandscape] is null, this will create a spec for both portrait and landscape, for
-         * each of the light/dark themes.
-         */
-        fun forDisplays(
-            vararg displays: DisplaySpec,
-            isDarkTheme: Boolean? = null,
-            isLandscape: Boolean? = null,
-        ): List<ScreenshotTestSpec> {
-            return displays.flatMap { display ->
-                buildList {
-                    fun addDisplay(isLandscape: Boolean) {
-                        if (isDarkTheme != true) {
-                            add(ScreenshotTestSpec(display, isDarkTheme = false, isLandscape))
-                        }
-
-                        if (isDarkTheme != false) {
-                            add(ScreenshotTestSpec(display, isDarkTheme = true, isLandscape))
-                        }
-                    }
-
-                    if (isLandscape != true) {
-                        addDisplay(isLandscape = false)
-                    }
-
-                    if (isLandscape != false) {
-                        addDisplay(isLandscape = true)
-                    }
-                }
-            }
-        }
-    }
-
-    override fun toString(): String = buildString {
-        // This string is appended to PNGs stored in the device, so let's keep it simple.
-        append(display.name)
-        if (isDarkTheme) append("_dark")
-        if (isLandscape) append("_landscape")
-    }
-}
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/SystemUIGoldenImagePathManager.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/SystemUIGoldenImagePathManager.kt
new file mode 100644
index 0000000..cbab0a7
--- /dev/null
+++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/SystemUIGoldenImagePathManager.kt
@@ -0,0 +1,42 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.testing.screenshot
+
+import androidx.test.platform.app.InstrumentationRegistry
+import platform.test.screenshot.GoldenImagePathManager
+import platform.test.screenshot.PathConfig
+
+/** A [GoldenImagePathManager] that should be used for all SystemUI screenshot tests. */
+class SystemUIGoldenImagePathManager(
+    pathConfig: PathConfig,
+) :
+    GoldenImagePathManager(
+        appContext = InstrumentationRegistry.getInstrumentation().context,
+        deviceLocalPath =
+            InstrumentationRegistry.getInstrumentation()
+                .targetContext
+                .filesDir
+                .absolutePath
+                .toString() + "/sysui_screenshots",
+        pathConfig = pathConfig,
+    ) {
+    override fun toString(): String {
+        // This string is appended to all actual/expected screenshots on the device, so make sure
+        // it is a static value.
+        return "SystemUIGoldenImagePathManager"
+    }
+}
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 6a80c48..3209c8b 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
@@ -1,3 +1,19 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.android.systemui.testing.screenshot
 
 import android.app.Activity
@@ -11,21 +27,35 @@
 import org.junit.rules.TestRule
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
+import platform.test.screenshot.DeviceEmulationRule
+import platform.test.screenshot.DeviceEmulationSpec
+import platform.test.screenshot.MaterialYouColorsRule
+import platform.test.screenshot.ScreenshotTestRule
+import platform.test.screenshot.getEmulatedDevicePathConfig
 
-/** A rule for View screenshot diff tests. */
-class ViewScreenshotTestRule(testSpec: ScreenshotTestSpec) : TestRule {
+/** A rule for View screenshot diff unit tests. */
+class ViewScreenshotTestRule(emulationSpec: DeviceEmulationSpec) : TestRule {
+    private val colorsRule = MaterialYouColorsRule()
+    private val deviceEmulationRule = DeviceEmulationRule(emulationSpec)
+    private val screenshotRule =
+        ScreenshotTestRule(
+            SystemUIGoldenImagePathManager(getEmulatedDevicePathConfig(emulationSpec))
+        )
     private val activityRule = ActivityScenarioRule(ScreenshotActivity::class.java)
-    private val screenshotRule = ScreenshotTestRule(testSpec)
-
-    private val delegate = RuleChain.outerRule(screenshotRule).around(activityRule)
+    private val delegateRule =
+        RuleChain.outerRule(colorsRule)
+            .around(deviceEmulationRule)
+            .around(screenshotRule)
+            .around(activityRule)
+    private val matcher = UnitTestBitmapMatcher
 
     override fun apply(base: Statement, description: Description): Statement {
-        return delegate.apply(base, description)
+        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 [testSpec].
+     * by [goldenIdentifier] in the context of [emulationSpec].
      */
     fun screenshotTest(
         goldenIdentifier: String,
@@ -46,13 +76,17 @@
             // Check that the content is what we expected.
             val content = activity.requireViewById<ViewGroup>(android.R.id.content)
             assertEquals(1, content.childCount)
-            screenshotRule.screenshotTest(goldenIdentifier, content.getChildAt(0))
+            screenshotRule.assertBitmapAgainstGolden(
+                content.getChildAt(0).drawIntoBitmap(),
+                goldenIdentifier,
+                matcher
+            )
         }
     }
 
     /**
      * Compare the content of the dialog provided by [dialogProvider] with the golden image
-     * identified by [goldenIdentifier] in the context of [testSpec].
+     * identified by [goldenIdentifier] in the context of [emulationSpec].
      */
     fun dialogScreenshotTest(
         goldenIdentifier: String,
@@ -81,7 +115,11 @@
             // Check that the content is what we expected.
             val dialog = dialog ?: error("dialog is null")
             try {
-                screenshotRule.screenshotTest(goldenIdentifier, dialog.window.decorView)
+                screenshotRule.assertBitmapAgainstGolden(
+                    dialog.window.decorView.drawIntoBitmap(),
+                    goldenIdentifier,
+                    matcher,
+                )
             } finally {
                 dialog.dismiss()
             }