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
index 97665145ce764..dedf0a7a742d6 100644
--- a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewCapture.kt
+++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewCapture.kt
@@ -53,6 +53,7 @@ import kotlinx.coroutines.runBlocking
 fun View.captureToBitmap(window: Window? = null): ListenableFuture<Bitmap> {
     val bitmapFuture: ResolvableFuture<Bitmap> = ResolvableFuture.create()
     val mainExecutor = HandlerExecutor(Handler(Looper.getMainLooper()))
+    val isRobolectric = if (Build.FINGERPRINT.contains("robolectric")) true else false
 
     // disable drawing again if necessary once work is complete
     if (!HardwareRendererCompat.isDrawingEnabled()) {
@@ -61,8 +62,12 @@ fun View.captureToBitmap(window: Window? = null): ListenableFuture<Bitmap> {
     }
 
     mainExecutor.execute {
-        val forceRedrawFuture = forceRedraw()
-        forceRedrawFuture.addListener({ generateBitmap(bitmapFuture, window) }, mainExecutor)
+        if (isRobolectric) {
+            generateBitmap(bitmapFuture)
+        } else {
+            val forceRedrawFuture = forceRedraw()
+            forceRedrawFuture.addListener({ generateBitmap(bitmapFuture, window) }, mainExecutor)
+        }
     }
 
     return 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 2d47356f2ce16..f96d1e3d7fef6 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
@@ -19,6 +19,7 @@ package com.android.systemui.testing.screenshot
 import android.app.Activity
 import android.app.Dialog
 import android.graphics.Bitmap
+import android.os.Build
 import android.view.View
 import android.view.ViewGroup
 import android.view.ViewGroup.LayoutParams
@@ -26,6 +27,7 @@ import android.view.ViewGroup.LayoutParams.MATCH_PARENT
 import android.view.ViewGroup.LayoutParams.WRAP_CONTENT
 import androidx.activity.ComponentActivity
 import androidx.test.ext.junit.rules.ActivityScenarioRule
+import java.util.concurrent.TimeUnit
 import org.junit.Assert.assertEquals
 import org.junit.rules.RuleChain
 import org.junit.rules.TestRule
@@ -54,14 +56,14 @@ open class ViewScreenshotTestRule(
             )
         )
     private val activityRule = ActivityScenarioRule(ScreenshotActivity::class.java)
-    private val delegateRule =
-        RuleChain.outerRule(colorsRule)
-            .around(deviceEmulationRule)
-            .around(screenshotRule)
-            .around(activityRule)
+    private val roboRule =
+        RuleChain.outerRule(deviceEmulationRule).around(screenshotRule).around(activityRule)
+    private val delegateRule = RuleChain.outerRule(colorsRule).around(roboRule)
+    private val isRobolectric = if (Build.FINGERPRINT.contains("robolectric")) true else false
 
     override fun apply(base: Statement, description: Description): Statement {
-        return delegateRule.apply(base, description)
+        val ruleToApply = if (isRobolectric) roboRule else delegateRule
+        return ruleToApply.apply(base, description)
     }
 
     protected fun takeScreenshot(
@@ -94,7 +96,12 @@ open class ViewScreenshotTestRule(
             contentView = content.getChildAt(0)
         }
 
-        return contentView?.toBitmap() ?: error("contentView is null")
+        return if (isRobolectric) {
+            contentView?.captureToBitmap()?.get(10, TimeUnit.SECONDS)
+                ?: error("timeout while trying to capture view to bitmap")
+        } else {
+            contentView?.toBitmap() ?: error("contentView is null")
+        }
     }
 
     /**