diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index f27f400..edd5011 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -34,9 +34,9 @@
 
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.tapl.TestHelpers;
+import com.android.launcher3.util.rule.FailureWatcher;
 import com.android.systemui.shared.system.QuickStepContract;
 
-import org.junit.Assert;
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
@@ -79,6 +79,14 @@
                 description.getAnnotation(NavigationModeSwitch.class) != null) {
             Mode mode = description.getAnnotation(NavigationModeSwitch.class).mode();
             return new Statement() {
+                private void assertTrue(String message, boolean condition) {
+                    if(!condition) {
+                        final AssertionError assertionError = new AssertionError(message);
+                        FailureWatcher.onError(mLauncher.getDevice(), description, assertionError);
+                        throw assertionError;
+                    }
+                }
+
                 @Override
                 public void evaluate() throws Throwable {
                     mLauncher.enableDebugTracing();
@@ -107,7 +115,8 @@
                         Log.e(TAG, "Exception", e);
                         throw e;
                     } finally {
-                        Assert.assertTrue(setActiveOverlay(prevOverlayPkg, originalMode));
+                        assertTrue("Couldn't set overlay",
+                                setActiveOverlay(prevOverlayPkg, originalMode));
                     }
                     mLauncher.disableDebugTracing();
                 }
@@ -176,7 +185,7 @@
                         latch.await(10, TimeUnit.SECONDS);
                         targetContext.getMainExecutor().execute(() ->
                                 sysUINavigationMode.removeModeChangeListener(listener));
-                        Assert.assertTrue("Navigation mode didn't change to " + expectedMode,
+                        assertTrue("Navigation mode didn't change to " + expectedMode,
                                 currentSysUiNavigationMode() == expectedMode);
                     }
 
@@ -184,7 +193,7 @@
                         if (mLauncher.getNavigationModel() == expectedMode) break;
                         Thread.sleep(100);
                     }
-                    Assert.assertTrue("Couldn't switch to " + overlayPackage,
+                    assertTrue("Couldn't switch to " + overlayPackage,
                             mLauncher.getNavigationModel() == expectedMode);
 
                     for (int i = 0; i != 100; ++i) {
@@ -192,7 +201,7 @@
                         Thread.sleep(100);
                     }
                     final String error = mLauncher.getNavigationModeMismatchError();
-                    Assert.assertTrue("Switching nav mode: " + error, error == null);
+                    assertTrue("Switching nav mode: " + error, error == null);
 
                     Thread.sleep(5000);
                     return true;
diff --git a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
index e7a2bca..cdda0f0 100644
--- a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
+++ b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
@@ -21,10 +21,10 @@
         mDevice = device;
     }
 
-    private void dumpViewHierarchy() {
+    private static void dumpViewHierarchy(UiDevice device) {
         final ByteArrayOutputStream stream = new ByteArrayOutputStream();
         try {
-            mDevice.dumpWindowHierarchy(stream);
+            device.dumpWindowHierarchy(stream);
             stream.flush();
             stream.close();
             for (String line : stream.toString().split("\\r?\\n")) {
@@ -37,7 +37,11 @@
 
     @Override
     protected void failed(Throwable e, Description description) {
-        if (mDevice == null) return;
+        onError(mDevice, description, e);
+    }
+
+    public static void onError(UiDevice device, Description description, Throwable e) {
+        if (device == null) return;
         final String pathname = getInstrumentation().getTargetContext().
                 getFilesDir().getPath() + "/TestScreenshot-" + description.getMethodName()
                 + ".png";
@@ -45,15 +49,15 @@
                 ", screenshot will be saved to " + pathname +
                 ", track trace is below, UI object dump is further below:\n" +
                 Log.getStackTraceString(e));
-        dumpViewHierarchy();
+        dumpViewHierarchy(device);
 
         try {
-            final String dumpsysResult = mDevice.executeShellCommand(
+            final String dumpsysResult = device.executeShellCommand(
                     "dumpsys activity service TouchInteractionService");
             Log.d(TAG, "TouchInteractionService: " + dumpsysResult);
         } catch (IOException ex) {
         }
 
-        mDevice.takeScreenshot(new File(pathname));
+        device.takeScreenshot(new File(pathname));
     }
 }