diff --git a/tests/InputMethodStressTest/AndroidTest.xml b/tests/InputMethodStressTest/AndroidTest.xml
index 5fb260f..9ac4135 100644
--- a/tests/InputMethodStressTest/AndroidTest.xml
+++ b/tests/InputMethodStressTest/AndroidTest.xml
@@ -31,4 +31,10 @@
     <test class="com.android.tradefed.testtype.AndroidJUnitTest">
         <option name="package" value="com.android.inputmethod.stresstest" />
     </test>
+
+    <!-- Collect the files in the dump directory for debugging -->
+    <metrics_collector class="com.android.tradefed.device.metric.FilePullerLogCollector">
+        <option name="directory-keys" value="/sdcard/InputMethodStressTest/" />
+        <option name="collect-on-run-ended-only" value="true" />
+    </metrics_collector>
 </configuration>
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
index f0f78740..c84c2bc 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
@@ -48,6 +48,10 @@
     @Rule
     public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
 
+    @Rule
+    public ScreenCaptureRule mScreenCaptureRule =
+            new ScreenCaptureRule("/sdcard/InputMethodStressTest");
+
     @Test
     public void autoShow() {
         Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
index 4be07fb..1c957d4d 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
@@ -55,6 +55,10 @@
     @Rule
     public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
 
+    @Rule
+    public ScreenCaptureRule mScreenCaptureRule =
+            new ScreenCaptureRule("/sdcard/InputMethodStressTest");
+
     @Test
     public void test() {
         Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
index 356c470..31f00c6 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
@@ -75,6 +75,10 @@
     @Rule
     public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
 
+    @Rule
+    public ScreenCaptureRule mScreenCaptureRule =
+            new ScreenCaptureRule("/sdcard/InputMethodStressTest");
+
     private Context mContext;
     private NotificationManager mNotificationManager;
     private UiDevice mUiDevice;
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ScreenCaptureRule.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ScreenCaptureRule.java
new file mode 100644
index 0000000..4e4ef2e
--- /dev/null
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ScreenCaptureRule.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.inputmethod.stresstest;
+
+import com.android.compatibility.common.util.SystemUtil;
+
+import org.junit.rules.TestWatcher;
+import org.junit.runner.Description;
+
+import java.text.SimpleDateFormat;
+import java.util.Date;
+
+/**
+ * Takes a screenshot when the test fails.
+ *
+ * <p>Use {@link com.android.tradefed.device.metric.FilePullerLogCollector} to collect screenshots
+ * taken.
+ * For example, in AndroidTest.xml:
+ * <code>
+ *     <metrics_collector class="com.android.tradefed.device.metric.FilePullerLogCollector">
+ *         <option name="directory-keys" value="/sdcard/MyTest/" />
+ *         <option name="collect-on-run-ended-only" value="true" />
+ *     </metrics_collector>
+ * </code>
+ * in MyTest.java:
+ * <code>
+ *     @Rule
+ *     public ScreenCaptureRule mScreenCaptureRule = new ScreenCaptureRule("/sdcard/MyTest");
+ * </code>
+ */
+public class ScreenCaptureRule extends TestWatcher {
+
+    private static final String TAG = "ScreenCaptureRule";
+
+    private final String mOutDir;
+
+    public ScreenCaptureRule(String outDir) {
+        mOutDir = outDir;
+    }
+
+    @Override
+    protected void failed(Throwable e, Description description) {
+        super.failed(e, description);
+        String time = new SimpleDateFormat("yyyyMMdd-HHmmss").format(new Date());
+        String fileName = "screenshot-" + time + ".png";
+        capture(fileName);
+    }
+
+    /** Take a screenshot. */
+    public void capture(String fileName) {
+        SystemUtil.runCommandAndPrintOnLogcat(TAG, String.format("mkdir -p %s", mOutDir));
+        SystemUtil.runCommandAndPrintOnLogcat(TAG,
+                String.format("screencap %s/%s", mOutDir, fileName));
+    }
+}