Moved LogcatOnFailureCollector's actual collection to onTestEnd().

Bug: 130441146
Change-Id: I3c8e385d3c7ba587f49fccf125810397554de3bc
(cherry picked from commit 294cb926ecc9701b6fad8672b2fa1ac95eae1982)

2 files changed, 63 insertions, 22 deletions

diff --git a/libraries/device-collectors/src/main/java/android/device/collectors/LogcatOnFailureCollector.java b/libraries/device-collectors/src/main/java/android/device/collectors/LogcatOnFailureCollector.java
index aaa9944dd..ba1292778 100644
--- a/libraries/device-collectors/src/main/java/android/device/collectors/LogcatOnFailureCollector.java
+++ b/libraries/device-collectors/src/main/java/android/device/collectors/LogcatOnFailureCollector.java
@@ -26,12 +26,11 @@ import org.junit.runner.notification.Failure;
 
 import java.io.File;
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.Date;
 import java.util.HashMap;
 import java.text.SimpleDateFormat;
 
-import java.util.Arrays;
-
 /**
  * A {@link BaseMetricListener} that captures logcat after each test case failure.
  *
@@ -45,12 +44,14 @@ public class LogcatOnFailureCollector extends BaseMetricListener {
     static final SimpleDateFormat DATE_FORMATTER = new SimpleDateFormat("MM-dd HH:mm:ss.SSS");
 
     @VisibleForTesting static final String METRIC_SEP = "-";
+    @VisibleForTesting static final String FILENAME_SUFFIX = "logcat";
 
     public static final String DEFAULT_DIR = "run_listeners/logcats";
     private static final int BUFFER_SIZE = 16 * 1024;
 
     private File mDestDir;
     private String mStartTime = null;
+    private boolean mTestFailed = false;
 
     // Map to keep track of test iterations for multiple test iterations.
     private HashMap<Description, Integer> mTestIterations = new HashMap<>();
@@ -86,26 +87,45 @@ public class LogcatOnFailureCollector extends BaseMetricListener {
         mTestIterations.computeIfAbsent(description, desc -> 1);
     }
 
+    /**
+     * Mark the test as failed if this is called. The actual collection will be done in {@link
+     * onTestEnd} to ensure that all actions around a test failure end up in the logcat.
+     */
     @Override
     public void onTestFail(DataRecord testData, Description description, Failure failure) {
-        // Capture logcat from start time
-        if (mDestDir == null) {
-            return;
-        }
-        try {
-            int iteration = mTestIterations.get(description);
-            final String fileName =
-                    String.format(
-                            "%s.%s%s.txt",
-                            description.getClassName(),
-                            description.getMethodName(),
-                            iteration == 1 ? "" : (METRIC_SEP + String.valueOf(iteration)));
-            File logcat = new File(mDestDir, fileName);
-            getLogcatSince(mStartTime, logcat);
-            testData.addFileMetric(String.format("%s_%s", getTag(), logcat.getName()), logcat);
-        } catch (IOException | InterruptedException e) {
-            Log.e(getTag(), "Error trying to retrieve logcat.", e);
+        mTestFailed = true;
+    }
+
+    /** If the test fails, collect logcat since test start time. */
+    @Override
+    public void onTestEnd(DataRecord testData, Description description) {
+        if (mTestFailed) {
+            // Capture logcat from start time
+            if (mDestDir == null) {
+                return;
+            }
+            try {
+                int iteration = mTestIterations.get(description);
+                final String fileName =
+                        String.format(
+                                "%s.%s%s%s.txt",
+                                description.getClassName(),
+                                description.getMethodName(),
+                                iteration == 1 ? "" : (METRIC_SEP + String.valueOf(iteration)),
+                                METRIC_SEP + FILENAME_SUFFIX);
+                File logcat = new File(mDestDir, fileName);
+                getLogcatSince(mStartTime, logcat);
+                testData.addFileMetric(String.format("%s_%s", getTag(), logcat.getName()), logcat);
+            } catch (IOException | InterruptedException e) {
+                Log.e(getTag(), "Error trying to retrieve logcat.", e);
+            }
         }
+        // Reset the flag here, as onTestStart might not have been called if a @BeforeClass method
+        // fails.
+        mTestFailed = false;
+        // Update the start time here in case onTestStart() is not called for the next test. If it
+        // is called, the start time will be overwritten.
+        mStartTime = getCurrentDate();
     }
 
     /** @hide */
diff --git a/libraries/device-collectors/src/test/java/android/device/collectors/LogcatOnFailureCollectorTest.java b/libraries/device-collectors/src/test/java/android/device/collectors/LogcatOnFailureCollectorTest.java
index 30069834b..4485f4cf7 100644
--- a/libraries/device-collectors/src/test/java/android/device/collectors/LogcatOnFailureCollectorTest.java
+++ b/libraries/device-collectors/src/test/java/android/device/collectors/LogcatOnFailureCollectorTest.java
@@ -142,7 +142,14 @@ public final class LogcatOnFailureCollectorTest {
         assertEquals(1, check.size());
         // The only key is ours
         for (String key: check.keySet()) {
-            assertTrue(key.contains("run.test.txt"));
+            assertTrue(
+                    key.contains(
+                            String.join(
+                                    "",
+                                    "run.test",
+                                    LogcatOnFailureCollector.METRIC_SEP
+                                            + LogcatOnFailureCollector.FILENAME_SUFFIX,
+                                    ".txt")));
         }
     }
 
@@ -341,17 +348,31 @@ public final class LogcatOnFailureCollectorTest {
         // The first bundle should have the first logcat file, for the first iteration.
         Bundle check1 = capturedBundles.get(0);
         assertEquals(1, check1.size());
+        String expectedKey1 =
+                String.join(
+                        "",
+                        "run.test",
+                        LogcatOnFailureCollector.METRIC_SEP
+                                + LogcatOnFailureCollector.FILENAME_SUFFIX,
+                        ".txt");
         for (String key : check1.keySet()) {
             // The first iteration should not have an iteration number.
-            assertTrue(key.contains("run.test.txt"));
+            assertTrue(key.contains(expectedKey1));
         }
 
         // The second bundle should have the second logcat file, for the third iteration.
         Bundle check2 = capturedBundles.get(1);
         assertEquals(1, check2.size());
+        String expectedKey2 =
+                String.join(
+                        "",
+                        "run.test-3",
+                        LogcatOnFailureCollector.METRIC_SEP
+                                + LogcatOnFailureCollector.FILENAME_SUFFIX,
+                        ".txt");
         for (String key : check2.keySet()) {
             // The third iteration should have an iteration number, 3.
-            assertTrue(key.contains("run.test-3.txt"));
+            assertTrue(key.contains(expectedKey2));
         }
     }
 }