Android: Make it more obvious when batched tests fail in @AfterTest

I was confused why all tests were failing even though when run
individually most pass.

Bug: 345269210
Change-Id: Idbb761565caee9f9d7971f27f4d851112bf58aa3
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5663678
Reviewed-by: Henrique Nakashima <hnakashima@chromium.org>
Auto-Submit: Andrew Grieve <agrieve@chromium.org>
Commit-Queue: Andrew Grieve <agrieve@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1320659}

diff --git a/base/test/android/javatests/src/org/chromium/base/test/BaseJUnit4ClassRunner.java b/base/test/android/javatests/src/org/chromium/base/test/BaseJUnit4ClassRunner.java
index fc5306c..9447ade 100644
--- a/base/test/android/javatests/src/org/chromium/base/test/BaseJUnit4ClassRunner.java
+++ b/base/test/android/javatests/src/org/chromium/base/test/BaseJUnit4ClassRunner.java

@@ -98,6 +98,30 @@
         public void run(Context targetContext, FrameworkMethod testMethod);
     }
 
+    /** Makes it more obvious that all tests are being marked as failed. */
+    private static class BeforeClassException extends RuntimeException {
+        private BeforeClassException(boolean batchedTest, Throwable causedBy) {
+            super(
+                    "Exception in @BeforeClass."
+                            + (batchedTest
+                                    ? " All tests in this class will be marked as failed."
+                                    : ""),
+                    causedBy);
+        }
+    }
+
+    /** Makes it more obvious that all tests are being marked as failed. */
+    private static class AfterClassException extends RuntimeException {
+        private AfterClassException(boolean batchedTest, Throwable causedBy) {
+            super(
+                    "Exception in @AfterClass."
+                            + (batchedTest
+                                    ? " All tests in this class will be marked as failed."
+                                    : ""),
+                    causedBy);
+        }
+    }
+
     /**
      * Thrown if a prior test in the same batch also failed thus possibly causing the current test's
      * failure.
@@ -319,17 +343,17 @@
         BaseJUnit4ClassRunner.getApplication().getSystemService(JobScheduler.class).cancelAll();
     }
 
-    private static boolean isInUnitTestBatch(FrameworkMethod testMethod) {
-        Batch annotation = testMethod.getDeclaringClass().getAnnotation(Batch.class);
+    private static boolean isInUnitTestBatch(Class<?> testClass) {
+        Batch annotation = testClass.getAnnotation(Batch.class);
         return annotation != null && annotation.value().equals(Batch.UNIT_TESTS);
     }
 
-    private static boolean isBatchedTest(FrameworkMethod testMethod) {
-        return testMethod.getDeclaringClass().getAnnotation(Batch.class) != null;
+    private static boolean isBatchedTest(Class<?> testClass) {
+        return testClass.getAnnotation(Batch.class) != null;
     }
 
     private static void blockUnitTestsFromStartingBrowser(FrameworkMethod testMethod) {
-        if (isInUnitTestBatch(testMethod)) {
+        if (isInUnitTestBatch(testMethod.getDeclaringClass())) {
             if (testMethod.getAnnotation(RequiresRestart.class) != null) return;
             LibraryLoader.setBrowserProcessStartupBlockedForTesting();
         }
@@ -341,7 +365,14 @@
         return new Statement() {
             @Override
             public void evaluate() throws Throwable {
-                onBeforeTestClass();
+                try {
+                    onBeforeTestClass();
+                } catch (Throwable t) {
+                    if (t instanceof AssumptionViolatedException) {
+                        throw t;
+                    }
+                    throw new BeforeClassException(isBatchedTest(getTestClass().getJavaClass()), t);
+                }
                 Throwable exception = null;
                 try {
                     innerStatement.evaluate();
@@ -359,7 +390,8 @@
                     }
                 }
                 if (exception != null) {
-                    throw exception;
+                    throw new AfterClassException(
+                            isBatchedTest(getTestClass().getJavaClass()), exception);
                 }
             }
         };
@@ -403,7 +435,10 @@
     }
 
     private void markBatchTestFailed(FrameworkMethod method) {
-        if (mFailedBatchTestName == null && isBatchedTest(method) && !isInUnitTestBatch(method)) {
+        Class<?> testClass = method.getDeclaringClass();
+        if (mFailedBatchTestName == null
+                && isBatchedTest(testClass)
+                && !isInUnitTestBatch(testClass)) {
             mFailedBatchTestName = method.getName();
         }
     }