diff --git a/packages/SystemUI/screenshot/Android.bp b/packages/SystemUI/screenshot/Android.bp
index 601e92f..f449398 100644
--- a/packages/SystemUI/screenshot/Android.bp
+++ b/packages/SystemUI/screenshot/Android.bp
@@ -38,6 +38,7 @@
         "androidx.test.espresso.core",
         "androidx.appcompat_appcompat",
         "platform-screenshot-diff-core",
+        "guava",
     ],
 
     kotlincflags: ["-Xjvm-default=all"],
diff --git a/packages/SystemUI/screenshot/res/values/themes.xml b/packages/SystemUI/screenshot/res/values/themes.xml
index 40e50bb..a7f8a26 100644
--- a/packages/SystemUI/screenshot/res/values/themes.xml
+++ b/packages/SystemUI/screenshot/res/values/themes.xml
@@ -19,6 +19,12 @@
         <item name="android:windowActionBar">false</item>
         <item name="android:windowNoTitle">true</item>
 
+        <!-- We make the status and navigation bars transparent so that the screenshotted content is
+             not clipped by the status bar height when drawn into the Bitmap (which is what happens
+             given that we draw the view into the Bitmap using hardware acceleration). -->
+        <item name="android:statusBarColor">@android:color/transparent</item>
+        <item name="android:navigationBarColor">@android:color/transparent</item>
+
         <!-- Make sure that device specific cutouts don't impact the outcome of screenshot tests -->
         <item name="android:windowLayoutInDisplayCutoutMode">shortEdges</item>
     </style>
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
index 3d26cda..a4a70a4 100644
--- a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/Bitmap.kt
+++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/Bitmap.kt
@@ -24,6 +24,8 @@
 import platform.test.screenshot.matchers.PixelPerfectMatcher
 
 /** Draw this [View] into a [Bitmap]. */
+// TODO(b/195673633): Remove this once Compose screenshot tests use hardware rendering for their
+// tests.
 fun View.drawIntoBitmap(): Bitmap {
     val bitmap =
         Bitmap.createBitmap(
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewCapture.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewCapture.kt
new file mode 100644
index 0000000..c609e6f
--- /dev/null
+++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewCapture.kt
@@ -0,0 +1,180 @@
+package com.android.systemui.testing.screenshot
+
+import android.app.Activity
+import android.content.Context
+import android.content.ContextWrapper
+import android.graphics.Bitmap
+import android.graphics.Canvas
+import android.graphics.Rect
+import android.os.Build
+import android.os.Handler
+import android.os.Looper
+import android.util.Log
+import android.view.PixelCopy
+import android.view.SurfaceView
+import android.view.View
+import android.view.ViewTreeObserver
+import android.view.Window
+import androidx.annotation.RequiresApi
+import androidx.concurrent.futures.ResolvableFuture
+import androidx.test.annotation.ExperimentalTestApi
+import androidx.test.core.internal.os.HandlerExecutor
+import androidx.test.platform.graphics.HardwareRendererCompat
+import com.google.common.util.concurrent.ListenableFuture
+
+/*
+ * This file was forked from androidx/test/core/view/ViewCapture.kt to add [Window] parameter to
+ * [View.captureToBitmap].
+ * TODO(b/195673633): Remove this fork and use the AndroidX version instead.
+ */
+
+/**
+ * Asynchronously captures an image of the underlying view into a [Bitmap].
+ *
+ * For devices below [Build.VERSION_CODES#O] (or if the view's window cannot be determined), the
+ * image is obtained using [View#draw]. Otherwise, [PixelCopy] is used.
+ *
+ * This method will also enable [HardwareRendererCompat#setDrawingEnabled(boolean)] if required.
+ *
+ * This API is primarily intended for use in lower layer libraries or frameworks. For test authors,
+ * its recommended to use espresso or compose's captureToImage.
+ *
+ * This API is currently experimental and subject to change or removal.
+ */
+@ExperimentalTestApi
+@RequiresApi(Build.VERSION_CODES.JELLY_BEAN)
+fun View.captureToBitmap(window: Window? = null): ListenableFuture<Bitmap> {
+    val bitmapFuture: ResolvableFuture<Bitmap> = ResolvableFuture.create()
+    val mainExecutor = HandlerExecutor(Handler(Looper.getMainLooper()))
+
+    // disable drawing again if necessary once work is complete
+    if (!HardwareRendererCompat.isDrawingEnabled()) {
+        HardwareRendererCompat.setDrawingEnabled(true)
+        bitmapFuture.addListener({ HardwareRendererCompat.setDrawingEnabled(false) }, mainExecutor)
+    }
+
+    mainExecutor.execute {
+        val forceRedrawFuture = forceRedraw()
+        forceRedrawFuture.addListener({ generateBitmap(bitmapFuture, window) }, mainExecutor)
+    }
+
+    return bitmapFuture
+}
+
+/**
+ * Trigger a redraw of the given view.
+ *
+ * Should only be called on UI thread.
+ *
+ * @return a [ListenableFuture] that will be complete once ui drawing is complete
+ */
+// NoClassDefFoundError occurs on API 15
+@RequiresApi(Build.VERSION_CODES.JELLY_BEAN)
+// @RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
+@ExperimentalTestApi
+fun View.forceRedraw(): ListenableFuture<Void> {
+    val future: ResolvableFuture<Void> = ResolvableFuture.create()
+
+    if (Build.VERSION.SDK_INT >= 29 && isHardwareAccelerated) {
+        viewTreeObserver.registerFrameCommitCallback() { future.set(null) }
+    } else {
+        viewTreeObserver.addOnDrawListener(
+            object : ViewTreeObserver.OnDrawListener {
+                var handled = false
+                override fun onDraw() {
+                    if (!handled) {
+                        handled = true
+                        future.set(null)
+                        // cannot remove on draw listener inside of onDraw
+                        Handler(Looper.getMainLooper()).post {
+                            viewTreeObserver.removeOnDrawListener(this)
+                        }
+                    }
+                }
+            }
+        )
+    }
+    invalidate()
+    return future
+}
+
+private fun View.generateBitmap(
+    bitmapFuture: ResolvableFuture<Bitmap>,
+    window: Window? = null,
+) {
+    if (bitmapFuture.isCancelled) {
+        return
+    }
+    val destBitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
+    when {
+        Build.VERSION.SDK_INT < 26 -> generateBitmapFromDraw(destBitmap, bitmapFuture)
+        this is SurfaceView -> generateBitmapFromSurfaceViewPixelCopy(destBitmap, bitmapFuture)
+        else -> {
+            val window = window ?: getActivity()?.window
+            if (window != null) {
+                generateBitmapFromPixelCopy(window, destBitmap, bitmapFuture)
+            } else {
+                Log.i(
+                    "View.captureToImage",
+                    "Could not find window for view. Falling back to View#draw instead of PixelCopy"
+                )
+                generateBitmapFromDraw(destBitmap, bitmapFuture)
+            }
+        }
+    }
+}
+
+@SuppressWarnings("NewApi")
+private fun SurfaceView.generateBitmapFromSurfaceViewPixelCopy(
+    destBitmap: Bitmap,
+    bitmapFuture: ResolvableFuture<Bitmap>
+) {
+    val onCopyFinished =
+        PixelCopy.OnPixelCopyFinishedListener { result ->
+            if (result == PixelCopy.SUCCESS) {
+                bitmapFuture.set(destBitmap)
+            } else {
+                bitmapFuture.setException(
+                    RuntimeException(String.format("PixelCopy failed: %d", result))
+                )
+            }
+        }
+    PixelCopy.request(this, null, destBitmap, onCopyFinished, handler)
+}
+
+internal fun View.generateBitmapFromDraw(
+    destBitmap: Bitmap,
+    bitmapFuture: ResolvableFuture<Bitmap>
+) {
+    destBitmap.density = resources.displayMetrics.densityDpi
+    computeScroll()
+    val canvas = Canvas(destBitmap)
+    canvas.translate((-scrollX).toFloat(), (-scrollY).toFloat())
+    draw(canvas)
+    bitmapFuture.set(destBitmap)
+}
+
+private fun View.getActivity(): Activity? {
+    fun Context.getActivity(): Activity? {
+        return when (this) {
+            is Activity -> this
+            is ContextWrapper -> this.baseContext.getActivity()
+            else -> null
+        }
+    }
+    return context.getActivity()
+}
+
+private fun View.generateBitmapFromPixelCopy(
+    window: Window,
+    destBitmap: Bitmap,
+    bitmapFuture: ResolvableFuture<Bitmap>
+) {
+    val locationInWindow = intArrayOf(0, 0)
+    getLocationInWindow(locationInWindow)
+    val x = locationInWindow[0]
+    val y = locationInWindow[1]
+    val boundsInWindow = Rect(x, y, x + width, y + height)
+
+    return window.generateBitmapFromPixelCopy(boundsInWindow, destBitmap, bitmapFuture)
+}
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 3209c8b..60130e1 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
@@ -18,10 +18,22 @@
 
 import android.app.Activity
 import android.app.Dialog
+import android.graphics.Bitmap
+import android.graphics.HardwareRenderer
+import android.os.Looper
 import android.view.View
 import android.view.ViewGroup
 import android.view.ViewGroup.LayoutParams
+import android.view.ViewGroup.LayoutParams.MATCH_PARENT
+import android.view.ViewGroup.LayoutParams.WRAP_CONTENT
+import android.view.Window
+import androidx.activity.ComponentActivity
+import androidx.test.espresso.Espresso
 import androidx.test.ext.junit.rules.ActivityScenarioRule
+import com.google.common.util.concurrent.FutureCallback
+import com.google.common.util.concurrent.Futures
+import kotlin.coroutines.suspendCoroutine
+import kotlinx.coroutines.runBlocking
 import org.junit.Assert.assertEquals
 import org.junit.rules.RuleChain
 import org.junit.rules.TestRule
@@ -59,29 +71,39 @@
      */
     fun screenshotTest(
         goldenIdentifier: String,
-        layoutParams: LayoutParams =
-            LayoutParams(LayoutParams.WRAP_CONTENT, LayoutParams.WRAP_CONTENT),
-        viewProvider: (Activity) -> View,
+        mode: Mode = Mode.WrapContent,
+        viewProvider: (ComponentActivity) -> 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(viewProvider(activity), layoutParams)
+            val window = activity.window
+            window.setDecorFitsSystemWindows(false)
+
+            // Set the content.
+            activity.setContentView(viewProvider(activity), mode.layoutParams)
+
+            // Elevation/shadows is not deterministic when doing hardware rendering, so we disable
+            // it for any view in the hierarchy.
+            window.decorView.removeElevationRecursively()
         }
 
         // We call onActivity again because it will make sure that our Activity is done measuring,
         // laying out and drawing its content (that we set in the previous onActivity lambda).
+        var contentView: View? = null
         activityRule.scenario.onActivity { activity ->
             // Check that the content is what we expected.
             val content = activity.requireViewById<ViewGroup>(android.R.id.content)
             assertEquals(1, content.childCount)
-            screenshotRule.assertBitmapAgainstGolden(
-                content.getChildAt(0).drawIntoBitmap(),
-                goldenIdentifier,
-                matcher
-            )
+            contentView = content.getChildAt(0)
         }
+
+        val bitmap = contentView?.toBitmap() ?: error("contentView is null")
+        screenshotRule.assertBitmapAgainstGolden(
+            bitmap,
+            goldenIdentifier,
+            matcher,
+        )
     }
 
     /**
@@ -104,25 +126,78 @@
                     create()
                     window.setWindowAnimations(0)
 
+                    // Elevation/shadows is not deterministic when doing hardware rendering, so we
+                    // disable it for any view in the hierarchy.
+                    window.decorView.removeElevationRecursively()
+
                     // Show the dialog.
                     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")
-            try {
-                screenshotRule.assertBitmapAgainstGolden(
-                    dialog.window.decorView.drawIntoBitmap(),
-                    goldenIdentifier,
-                    matcher,
+        try {
+            val bitmap = dialog?.toBitmap() ?: error("dialog is null")
+            screenshotRule.assertBitmapAgainstGolden(
+                bitmap,
+                goldenIdentifier,
+                matcher,
+            )
+        } finally {
+            dialog?.dismiss()
+        }
+    }
+
+    private fun View.removeElevationRecursively() {
+        this.elevation = 0f
+
+        if (this is ViewGroup) {
+            repeat(childCount) { i -> getChildAt(i).removeElevationRecursively() }
+        }
+    }
+
+    private fun Dialog.toBitmap(): Bitmap {
+        val window = window
+        return window.decorView.toBitmap(window)
+    }
+
+    private fun View.toBitmap(window: Window? = null): Bitmap {
+        if (Looper.getMainLooper() == Looper.myLooper()) {
+            error("toBitmap() can't be called from the main thread")
+        }
+
+        if (!HardwareRenderer.isDrawingEnabled()) {
+            error("Hardware rendering is not enabled")
+        }
+
+        // Make sure we are idle.
+        Espresso.onIdle()
+
+        val mainExecutor = context.mainExecutor
+        return runBlocking {
+            suspendCoroutine { continuation ->
+                Futures.addCallback(
+                    captureToBitmap(window),
+                    object : FutureCallback<Bitmap> {
+                        override fun onSuccess(result: Bitmap?) {
+                            continuation.resumeWith(Result.success(result!!))
+                        }
+
+                        override fun onFailure(t: Throwable) {
+                            continuation.resumeWith(Result.failure(t))
+                        }
+                    },
+                    // We know that we are not on the main thread, so we can block the current
+                    // thread and wait for the result in the main thread.
+                    mainExecutor,
                 )
-            } finally {
-                dialog.dismiss()
             }
         }
     }
+
+    enum class Mode(val layoutParams: LayoutParams) {
+        WrapContent(LayoutParams(WRAP_CONTENT, WRAP_CONTENT)),
+        MatchSize(LayoutParams(MATCH_PARENT, MATCH_PARENT)),
+        MatchWidth(LayoutParams(MATCH_PARENT, WRAP_CONTENT)),
+        MatchHeight(LayoutParams(WRAP_CONTENT, MATCH_PARENT)),
+    }
 }
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/WindowCapture.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/WindowCapture.kt
new file mode 100644
index 0000000..d34f46b
--- /dev/null
+++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/WindowCapture.kt
@@ -0,0 +1,37 @@
+package com.android.systemui.testing.screenshot
+
+import android.graphics.Bitmap
+import android.graphics.Rect
+import android.os.Handler
+import android.os.Looper
+import android.view.PixelCopy
+import android.view.Window
+import androidx.concurrent.futures.ResolvableFuture
+
+/*
+ * This file was forked from androidx/test/core/view/WindowCapture.kt.
+ * TODO(b/195673633): Remove this fork and use the AndroidX version instead.
+ */
+fun Window.generateBitmapFromPixelCopy(
+    boundsInWindow: Rect? = null,
+    destBitmap: Bitmap,
+    bitmapFuture: ResolvableFuture<Bitmap>
+) {
+    val onCopyFinished =
+        PixelCopy.OnPixelCopyFinishedListener { result ->
+            if (result == PixelCopy.SUCCESS) {
+                bitmapFuture.set(destBitmap)
+            } else {
+                bitmapFuture.setException(
+                    RuntimeException(String.format("PixelCopy failed: %d", result))
+                )
+            }
+        }
+    PixelCopy.request(
+        this,
+        boundsInWindow,
+        destBitmap,
+        onCopyFinished,
+        Handler(Looper.getMainLooper())
+    )
+}