diff --git a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
index 3edda6ba5b..7492ab8c16 100644
--- a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
@@ -122,7 +122,7 @@ public class FallbackRecentsTest {
         mOrderSensitiveRules = RuleChain
                 .outerRule(new SamplerRule())
                 .around(new NavigationModeSwitchRule(mLauncher))
-                .around(new FailureWatcher(mDevice, mLauncher, viewCaptureRule.getViewCapture()))
+                .around(new FailureWatcher(mLauncher, viewCaptureRule::getViewCaptureData))
                 .around(viewCaptureRule)
                 .around(new ViewCaptureAnalysisRule(viewCaptureRule.getViewCapture()));
 
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 68e586c5af..1262a26338 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -207,7 +207,7 @@ public abstract class AbstractLauncherUiTest {
         final ViewCaptureRule viewCaptureRule = new ViewCaptureRule(mActivityMonitor::getActivity);
         final RuleChain inner = RuleChain
                 .outerRule(new PortraitLandscapeRunner(this))
-                .around(new FailureWatcher(mDevice, mLauncher, viewCaptureRule.getViewCapture()))
+                .around(new FailureWatcher(mLauncher, viewCaptureRule::getViewCaptureData))
                 .around(viewCaptureRule)
                 .around(new ViewCaptureAnalysisRule(viewCaptureRule.getViewCapture()));
 
diff --git a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
index 7ca6a06ed2..f2ae9d3ef5 100644
--- a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
+++ b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
@@ -8,10 +8,9 @@ import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
-import androidx.test.core.app.ApplicationProvider;
 import androidx.test.uiautomator.UiDevice;
 
-import com.android.app.viewcapture.ViewCapture;
+import com.android.app.viewcapture.data.ExportedData;
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 
@@ -24,22 +23,21 @@ import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.util.function.Supplier;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipOutputStream;
 
 public class FailureWatcher extends TestWatcher {
     private static final String TAG = "FailureWatcher";
     private static boolean sSavedBugreport = false;
-    final private UiDevice mDevice;
     private final LauncherInstrumentation mLauncher;
     @NonNull
-    private final ViewCapture mViewCapture;
+    private final Supplier<ExportedData> mViewCaptureDataSupplier;
 
-    public FailureWatcher(UiDevice device, LauncherInstrumentation launcher,
-            @NonNull ViewCapture viewCapture) {
-        mDevice = device;
+    public FailureWatcher(LauncherInstrumentation launcher,
+            @NonNull Supplier<ExportedData> viewCaptureDataSupplier) {
         mLauncher = launcher;
-        mViewCapture = viewCapture;
+        mViewCaptureDataSupplier = viewCaptureDataSupplier;
     }
 
     @Override
@@ -71,7 +69,7 @@ public class FailureWatcher extends TestWatcher {
 
     @Override
     protected void failed(Throwable e, Description description) {
-        onError(mLauncher, description, e, mViewCapture);
+        onError(mLauncher, description, e, mViewCaptureDataSupplier);
     }
 
     static File diagFile(Description description, String prefix, String ext) {
@@ -86,7 +84,7 @@ public class FailureWatcher extends TestWatcher {
     }
 
     private static void onError(LauncherInstrumentation launcher, Description description,
-            Throwable e, @Nullable ViewCapture viewCapture) {
+            Throwable e, @Nullable Supplier<ExportedData> viewCaptureDataSupplier) {
 
         final File sceenshot = diagFile(description, "TestScreenshot", "png");
         final File hierarchy = diagFile(description, "Hierarchy", "zip");
@@ -103,9 +101,9 @@ public class FailureWatcher extends TestWatcher {
             dumpCommand("cmd window dump-visible-window-views", out);
             out.closeEntry();
 
-            if (viewCapture != null) {
+            if (viewCaptureDataSupplier != null) {
                 out.putNextEntry(new ZipEntry("FS/data/misc/wmtrace/failed_test.vc"));
-                viewCapture.dumpTo(out, ApplicationProvider.getApplicationContext());
+                viewCaptureDataSupplier.get().writeTo(out);
                 out.closeEntry();
             }
         } catch (Exception ignored) {
diff --git a/tests/src/com/android/launcher3/util/rule/ViewCaptureRule.kt b/tests/src/com/android/launcher3/util/rule/ViewCaptureRule.kt
index ca3147d811..6c065026b4 100644
--- a/tests/src/com/android/launcher3/util/rule/ViewCaptureRule.kt
+++ b/tests/src/com/android/launcher3/util/rule/ViewCaptureRule.kt
@@ -22,6 +22,7 @@ import android.os.Bundle
 import androidx.test.core.app.ApplicationProvider
 import com.android.app.viewcapture.SimpleViewCapture
 import com.android.app.viewcapture.ViewCapture.MAIN_EXECUTOR
+import com.android.app.viewcapture.data.ExportedData
 import com.android.launcher3.util.ActivityLifecycleCallbacksAdapter
 import java.util.function.Supplier
 import org.junit.rules.TestRule
@@ -36,10 +37,13 @@ import org.junit.runners.model.Statement
  */
 class ViewCaptureRule(var alreadyOpenActivitySupplier: Supplier<Activity?>) : TestRule {
     val viewCapture = SimpleViewCapture("test-view-capture")
+    var viewCaptureData: ExportedData? = null
+        private set
 
     override fun apply(base: Statement, description: Description): Statement {
         return object : Statement() {
             override fun evaluate() {
+                viewCaptureData = null
                 val windowListenerCloseables = mutableListOf<SafeCloseable>()
 
                 val alreadyOpenActivity = alreadyOpenActivitySupplier.get()
@@ -68,6 +72,9 @@ class ViewCaptureRule(var alreadyOpenActivitySupplier: Supplier<Activity?>) : Te
                 } finally {
                     application.unregisterActivityLifecycleCallbacks(lifecycleCallbacks)
 
+                    viewCaptureData =
+                        viewCapture.getExportedData(ApplicationProvider.getApplicationContext())
+
                     // Clean up ViewCapture references here rather than in onActivityDestroyed so
                     // test code can access view hierarchy capture. onActivityDestroyed would delete
                     // view capture data before FailureWatcher could output it as a test artifact.