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 2c3ff2c75c727..35812e39677d2 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,6 +1,7 @@
 package com.android.systemui.testing.screenshot
 
 import android.app.Activity
+import android.app.Dialog
 import android.view.View
 import android.view.ViewGroup
 import android.view.ViewGroup.LayoutParams
@@ -23,20 +24,20 @@ class ViewScreenshotTestRule(testSpec: ScreenshotTestSpec) : TestRule {
     }
 
     /**
-     * Compare the content of [view] with the golden image identified by [goldenIdentifier] in the
-     * context of [testSpec].
+     * Compare the content of the view provided by [viewProvider] with the golden image identified
+     * by [goldenIdentifier] in the context of [testSpec].
      */
     fun screenshotTest(
         goldenIdentifier: String,
         layoutParams: LayoutParams =
             LayoutParams(LayoutParams.WRAP_CONTENT, LayoutParams.WRAP_CONTENT),
-        view: (Activity) -> View,
+        viewProvider: (Activity) -> View,
     ) {
         activityRule.scenario.onActivity { activity ->
             // Make sure that the activity draws full screen and fits the whole display instead of
             // the system bars.
             activity.window.setDecorFitsSystemWindows(false)
-            activity.setContentView(view(activity), layoutParams)
+            activity.setContentView(viewProvider(activity), layoutParams)
         }
 
         // We call onActivity again because it will make sure that our Activity is done measuring,
@@ -48,4 +49,32 @@ class ViewScreenshotTestRule(testSpec: ScreenshotTestSpec) : TestRule {
             screenshotRule.screenshotTest(goldenIdentifier, content.getChildAt(0))
         }
     }
+
+    /**
+     * Compare the content of the dialog provided by [dialogProvider] with the golden image
+     * identified by [goldenIdentifier] in the context of [testSpec].
+     */
+    fun dialogScreenshotTest(
+        goldenIdentifier: String,
+        dialogProvider: (Activity) -> Dialog,
+    ) {
+        var dialog: Dialog? = null
+        activityRule.scenario.onActivity { activity ->
+            // Make sure that the dialog draws full screen and fits the whole display instead of the
+            // system bars.
+            dialog =
+                dialogProvider(activity).apply {
+                    window.setDecorFitsSystemWindows(false)
+                    show()
+                }
+        }
+
+        // We call onActivity again because it will make sure that our Dialog is done measuring,
+        // laying out and drawing its content (that we set in the previous onActivity lambda).
+        activityRule.scenario.onActivity {
+            // Check that the content is what we expected.
+            val dialog = dialog ?: error("dialog is null")
+            screenshotRule.screenshotTest(goldenIdentifier, dialog.window.decorView)
+        }
+    }
 }