Merge "[Ravenwood] Cleanup RATR implementation and project structure" into main

6 files changed, 722 insertions, 219 deletions

diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java
new file mode 100644
index 000000000000..30a653d2da76
--- /dev/null
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java
@@ -0,0 +1,453 @@
+/*
+ * Copyright (C) 2024 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 android.platform.test.ravenwood;
+
+import static com.android.ravenwood.common.RavenwoodCommonUtils.RAVENWOOD_VERBOSE_LOGGING;
+import static com.android.ravenwood.common.RavenwoodCommonUtils.ensureIsPublicVoidMethod;
+
+import static org.junit.Assert.fail;
+import static org.junit.Assume.assumeTrue;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.Bundle;
+import android.platform.test.annotations.RavenwoodTestRunnerInitializing;
+import android.platform.test.annotations.internal.InnerRunner;
+import android.platform.test.ravenwood.RavenwoodTestStats.Result;
+import android.util.Log;
+
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import org.junit.AssumptionViolatedException;
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runner.Runner;
+import org.junit.runner.manipulation.Filter;
+import org.junit.runner.manipulation.Filterable;
+import org.junit.runner.manipulation.NoTestsRemainException;
+import org.junit.runner.notification.Failure;
+import org.junit.runner.notification.RunNotifier;
+import org.junit.runners.BlockJUnit4ClassRunner;
+import org.junit.runners.Suite;
+import org.junit.runners.model.Statement;
+import org.junit.runners.model.TestClass;
+
+import java.lang.annotation.Annotation;
+import java.lang.reflect.InvocationTargetException;
+import java.util.function.BiConsumer;
+
+/**
+ * A test runner used for Ravenwood.
+ *
+ * It will delegate to another runner specified with {@link InnerRunner}
+ * (default = {@link BlockJUnit4ClassRunner}) with the following features.
+ * - Add a called before the inner runner gets a chance to run. This can be used to initialize
+ *   stuff used by the inner runner.
+ * - Add hook points with help from the four test rules such as {@link #sImplicitClassOuterRule},
+ *   which are also injected by the ravenizer tool.
+ *
+ * We use this runner to:
+ * - Initialize the Ravenwood environment.
+ * - Handle {@link android.platform.test.annotations.DisabledOnRavenwood}.
+ */
+public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase {
+    public static final String TAG = "Ravenwood";
+
+    /** Scope of a hook. */
+    public enum Scope {
+        Class,
+        Instance,
+    }
+
+    /** Order of a hook. */
+    public enum Order {
+        Outer,
+        Inner,
+    }
+
+    private record HookRule(Scope scope, Order order) implements TestRule {
+        @Override
+        public Statement apply(Statement base, Description description) {
+            return getCurrentRunner().wrapWithHooks(base, description, scope, order);
+        }
+    }
+
+    // The following four rule instances will be injected to tests by the Ravenizer tool.
+    public static final TestRule sImplicitClassOuterRule = new HookRule(Scope.Class, Order.Outer);
+    public static final TestRule sImplicitClassInnerRule = new HookRule(Scope.Class, Order.Inner);
+    public static final TestRule sImplicitInstOuterRule = new HookRule(Scope.Instance, Order.Outer);
+    public static final TestRule sImplicitInstInnerRule = new HookRule(Scope.Instance, Order.Inner);
+
+    /** Keeps track of the runner on the current thread. */
+    private static final ThreadLocal<RavenwoodAwareTestRunner> sCurrentRunner = new ThreadLocal<>();
+
+    static RavenwoodAwareTestRunner getCurrentRunner() {
+        var runner = sCurrentRunner.get();
+        if (runner == null) {
+            throw new RuntimeException("Current test runner not set!");
+        }
+        return runner;
+    }
+
+    private final Class<?> mTestJavaClass;
+    private TestClass mTestClass = null;
+    private Runner mRealRunner = null;
+    private Description mDescription = null;
+    private Throwable mExceptionInConstructor = null;
+
+    /**
+     * Stores internal states / methods associated with this runner that's only needed in
+     * junit-impl.
+     */
+    final RavenwoodRunnerState mState = new RavenwoodRunnerState(this);
+
+    public TestClass getTestClass() {
+        return mTestClass;
+    }
+
+    /**
+     * Constructor.
+     */
+    public RavenwoodAwareTestRunner(Class<?> testClass) {
+        RavenwoodRuntimeEnvironmentController.globalInitOnce();
+        mTestJavaClass = testClass;
+        try {
+            /*
+             * If the class has @DisabledOnRavenwood, then we'll delegate to
+             * ClassSkippingTestRunner, which simply skips it.
+             *
+             * We need to do it before instantiating TestClass for b/367694651.
+             */
+            if (!RavenwoodEnablementChecker.shouldRunClassOnRavenwood(testClass, true)) {
+                mRealRunner = new ClassSkippingTestRunner(testClass);
+                mDescription = mRealRunner.getDescription();
+                return;
+            }
+
+            mTestClass = new TestClass(testClass);
+
+            Log.v(TAG, "RavenwoodAwareTestRunner starting for " + testClass.getCanonicalName());
+
+            onRunnerInitializing();
+
+            mRealRunner = instantiateRealRunner(mTestClass);
+            mDescription = mRealRunner.getDescription();
+        } catch (Throwable th) {
+            // If we throw in the constructor, Tradefed may not report it and just ignore the class,
+            // so record it and throw it when the test actually started.
+            Log.e(TAG, "Fatal: Exception detected in constructor", th);
+            mExceptionInConstructor = new RuntimeException("Exception detected in constructor",
+                    th);
+            mDescription = Description.createTestDescription(testClass, "Constructor");
+
+            // This is for testing if tradefed is fixed.
+            if ("1".equals(System.getenv("RAVENWOOD_THROW_EXCEPTION_IN_TEST_RUNNER"))) {
+                throw th;
+            }
+        }
+    }
+
+    @Override
+    Runner getRealRunner() {
+        return mRealRunner;
+    }
+
+    /**
+     * Run the bare minimum setup to initialize the wrapped runner.
+     */
+    // This method is called by the ctor, so never make it virtual.
+    private void onRunnerInitializing() {
+        // This is needed to make AndroidJUnit4ClassRunner happy.
+        InstrumentationRegistry.registerInstance(null, Bundle.EMPTY);
+
+        // Hook point to allow more customization.
+        runAnnotatedMethodsOnRavenwood(RavenwoodTestRunnerInitializing.class, null);
+    }
+
+    private void runAnnotatedMethodsOnRavenwood(Class<? extends Annotation> annotationClass,
+            Object instance) {
+        Log.v(TAG, "runAnnotatedMethodsOnRavenwood() " + annotationClass.getName());
+
+        for (var method : getTestClass().getAnnotatedMethods(annotationClass)) {
+            ensureIsPublicVoidMethod(method.getMethod(), /* isStatic=*/ instance == null);
+
+            var methodDesc = method.getDeclaringClass().getName() + "."
+                    + method.getMethod().toString();
+            try {
+                method.getMethod().invoke(instance);
+            } catch (IllegalAccessException | InvocationTargetException e) {
+                throw logAndFail("Caught exception while running method " + methodDesc, e);
+            }
+        }
+    }
+
+    @Override
+    public Description getDescription() {
+        return mDescription;
+    }
+
+    @Override
+    public void run(RunNotifier realNotifier) {
+        final var notifier = new RavenwoodRunNotifier(realNotifier);
+        final var description = getDescription();
+
+        if (mRealRunner instanceof ClassSkippingTestRunner) {
+            mRealRunner.run(notifier);
+            Log.i(TAG, "onClassSkipped: description=" + description);
+            RavenwoodTestStats.getInstance().onClassSkipped(description);
+            return;
+        }
+
+        Log.v(TAG, "Starting " + mTestJavaClass.getCanonicalName());
+        if (RAVENWOOD_VERBOSE_LOGGING) {
+            dumpDescription(description);
+        }
+
+        if (maybeReportExceptionFromConstructor(notifier)) {
+            return;
+        }
+
+        // TODO(b/365976974): handle nested classes better
+        final boolean skipRunnerHook =
+                mRealRunnerTakesRunnerBuilder && mRealRunner instanceof Suite;
+
+        sCurrentRunner.set(this);
+        try {
+            if (!skipRunnerHook) {
+                try {
+                    mState.enterTestClass(description);
+                } catch (Throwable th) {
+                    notifier.reportBeforeTestFailure(description, th);
+                    return;
+                }
+            }
+
+            // Delegate to the inner runner.
+            mRealRunner.run(notifier);
+        } finally {
+            sCurrentRunner.remove();
+
+            if (!skipRunnerHook) {
+                try {
+                    RavenwoodTestStats.getInstance().onClassFinished(description);
+                    mState.exitTestClass();
+                } catch (Throwable th) {
+                    notifier.reportAfterTestFailure(th);
+                }
+            }
+        }
+    }
+
+    /** Throw the exception detected in the constructor, if any. */
+    private boolean maybeReportExceptionFromConstructor(RunNotifier notifier) {
+        if (mExceptionInConstructor == null) {
+            return false;
+        }
+        notifier.fireTestStarted(mDescription);
+        notifier.fireTestFailure(new Failure(mDescription, mExceptionInConstructor));
+        notifier.fireTestFinished(mDescription);
+
+        return true;
+    }
+
+    private Statement wrapWithHooks(Statement base, Description description, Scope scope,
+            Order order) {
+        return new Statement() {
+            @Override
+            public void evaluate() throws Throwable {
+                runWithHooks(description, scope, order, base);
+            }
+        };
+    }
+
+    private void runWithHooks(Description description, Scope scope, Order order, Statement s)
+            throws Throwable {
+        assumeTrue(onBefore(description, scope, order));
+        try {
+            s.evaluate();
+            onAfter(description, scope, order, null);
+        } catch (Throwable t) {
+            if (onAfter(description, scope, order, t)) {
+                throw t;
+            }
+        }
+    }
+
+    /**
+     * A runner that simply skips a class. It still has to support {@link Filterable}
+     * because otherwise the result still says "SKIPPED" even when it's not included in the
+     * filter.
+     */
+    private static class ClassSkippingTestRunner extends Runner implements Filterable {
+        private final Description mDescription;
+        private boolean mFilteredOut;
+
+        ClassSkippingTestRunner(Class<?> testClass) {
+            mDescription = Description.createTestDescription(testClass, testClass.getSimpleName());
+            mFilteredOut = false;
+        }
+
+        @Override
+        public Description getDescription() {
+            return mDescription;
+        }
+
+        @Override
+        public void run(RunNotifier notifier) {
+            if (mFilteredOut) {
+                return;
+            }
+            notifier.fireTestSuiteStarted(mDescription);
+            notifier.fireTestIgnored(mDescription);
+            notifier.fireTestSuiteFinished(mDescription);
+        }
+
+        @Override
+        public void filter(Filter filter) throws NoTestsRemainException {
+            if (filter.shouldRun(mDescription)) {
+                mFilteredOut = false;
+            } else {
+                throw new NoTestsRemainException();
+            }
+        }
+    }
+
+    /**
+     * Called before a test / class.
+     *
+     * Return false if it should be skipped.
+     */
+    private boolean onBefore(Description description, Scope scope, Order order) {
+        Log.v(TAG, "onBefore: description=" + description + ", " + scope + ", " + order);
+
+        if (scope == Scope.Instance && order == Order.Outer) {
+            // Start of a test method.
+            mState.enterTestMethod(description);
+        }
+
+        final var classDescription = mState.getClassDescription();
+
+        // Class-level annotations are checked by the runner already, so we only check
+        // method-level annotations here.
+        if (scope == Scope.Instance && order == Order.Outer) {
+            if (!RavenwoodEnablementChecker.shouldEnableOnRavenwood(description, true)) {
+                RavenwoodTestStats.getInstance().onTestFinished(
+                        classDescription, description, Result.Skipped);
+                return false;
+            }
+        }
+        return true;
+    }
+
+    /**
+     * Called after a test / class.
+     *
+     * Return false if the exception should be ignored.
+     */
+    private boolean onAfter(Description description, Scope scope, Order order, Throwable th) {
+        Log.v(TAG, "onAfter: description=" + description + ", " + scope + ", " + order + ", " + th);
+
+        final var classDescription = mState.getClassDescription();
+
+        if (scope == Scope.Instance && order == Order.Outer) {
+            // End of a test method.
+            mState.exitTestMethod();
+
+            final Result result;
+            if (th == null) {
+                result = Result.Passed;
+            } else if (th instanceof AssumptionViolatedException) {
+                result = Result.Skipped;
+            } else {
+                result = Result.Failed;
+            }
+
+            RavenwoodTestStats.getInstance().onTestFinished(classDescription, description, result);
+        }
+
+        // If RUN_DISABLED_TESTS is set, and the method did _not_ throw, make it an error.
+        if (RavenwoodRule.private$ravenwood().isRunningDisabledTests()
+                && scope == Scope.Instance && order == Order.Outer) {
+
+            boolean isTestEnabled = RavenwoodEnablementChecker.shouldEnableOnRavenwood(
+                    description, false);
+            if (th == null) {
+                // Test passed. Is the test method supposed to be enabled?
+                if (isTestEnabled) {
+                    // Enabled and didn't throw, okay.
+                    return true;
+                } else {
+                    // Disabled and didn't throw. We should report it.
+                    fail("Test wasn't included under Ravenwood, but it actually "
+                            + "passed under Ravenwood; consider updating annotations");
+                    return true; // unreachable.
+                }
+            } else {
+                // Test failed.
+                if (isTestEnabled) {
+                    // Enabled but failed. We should throw the exception.
+                    return true;
+                } else {
+                    // Disabled and failed. Expected. Don't throw.
+                    return false;
+                }
+            }
+        }
+        return true;
+    }
+
+    /**
+     * Called by RavenwoodRule.
+     */
+    static void onRavenwoodRuleEnter(Description description, RavenwoodRule rule) {
+        Log.v(TAG, "onRavenwoodRuleEnter: description=" + description);
+        getCurrentRunner().mState.enterRavenwoodRule(rule);
+    }
+
+    /**
+     * Called by RavenwoodRule.
+     */
+    static void onRavenwoodRuleExit(Description description, RavenwoodRule rule) {
+        Log.v(TAG, "onRavenwoodRuleExit: description=" + description);
+        getCurrentRunner().mState.exitRavenwoodRule(rule);
+    }
+
+    private void dumpDescription(Description desc) {
+        dumpDescription(desc, "[TestDescription]=", "  ");
+    }
+
+    private void dumpDescription(Description desc, String header, String indent) {
+        Log.v(TAG, indent + header + desc);
+
+        var children = desc.getChildren();
+        var childrenIndent = "  " + indent;
+        for (int i = 0; i < children.size(); i++) {
+            dumpDescription(children.get(i), "#" + i + ": ", childrenIndent);
+        }
+    }
+
+    static volatile BiConsumer<String, Throwable> sCriticalErrorHandler = null;
+
+    static void onCriticalError(@NonNull String message, @Nullable Throwable th) {
+        Log.e(TAG, "Critical error! " + message, th);
+        var handler = sCriticalErrorHandler;
+        if (handler == null) {
+            Log.e(TAG, "Ravenwood cannot continue. Killing self process.", th);
+            System.exit(1);
+        }
+        handler.accept(message, th);
+    }
+}
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
deleted file mode 100644
index e0f9ec94a819..000000000000
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
+++ /dev/null
@@ -1,209 +0,0 @@
-/*
- * Copyright (C) 2024 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 android.platform.test.ravenwood;
-
-import static org.junit.Assert.fail;
-
-import android.os.Bundle;
-import android.platform.test.ravenwood.RavenwoodAwareTestRunner.Order;
-import android.platform.test.ravenwood.RavenwoodAwareTestRunner.Scope;
-import android.platform.test.ravenwood.RavenwoodTestStats.Result;
-import android.util.Log;
-
-import androidx.test.platform.app.InstrumentationRegistry;
-
-import org.junit.AssumptionViolatedException;
-import org.junit.runner.Description;
-import org.junit.runners.model.TestClass;
-
-/**
- * Provide hook points created by {@link RavenwoodAwareTestRunner}.
- *
- * States are associated with each {@link RavenwoodAwareTestRunner} are stored in
- * {@link RavenwoodRunnerState}, rather than as members of {@link RavenwoodAwareTestRunner}.
- * See its javadoc for the reasons.
- *
- * All methods in this class must be called from the test main thread.
- */
-public class RavenwoodAwareTestRunnerHook {
-    private static final String TAG = RavenwoodAwareTestRunner.TAG;
-
-    private RavenwoodAwareTestRunnerHook() {
-    }
-
-    /**
-     * Called before any code starts. Internally it will only initialize the environment once.
-     */
-    public static void performGlobalInitialization() {
-        RavenwoodRuntimeEnvironmentController.globalInitOnce();
-    }
-
-    /**
-     * Called when a runner starts, before the inner runner gets a chance to run.
-     */
-    public static void onRunnerInitializing(RavenwoodAwareTestRunner runner, TestClass testClass) {
-        Log.i(TAG, "onRunnerInitializing: testClass=" + testClass.getJavaClass()
-                + " runner=" + runner);
-
-        // This is needed to make AndroidJUnit4ClassRunner happy.
-        InstrumentationRegistry.registerInstance(null, Bundle.EMPTY);
-    }
-
-    /**
-     * Called when a whole test class is skipped.
-     */
-    public static void onClassSkipped(Description description) {
-        Log.i(TAG, "onClassSkipped: description=" + description);
-        RavenwoodTestStats.getInstance().onClassSkipped(description);
-    }
-
-    /**
-     * Called before the inner runner starts.
-     */
-    public static void onBeforeInnerRunnerStart(
-            RavenwoodAwareTestRunner runner, Description description) throws Throwable {
-        Log.v(TAG, "onBeforeInnerRunnerStart: description=" + description);
-
-        // Prepare the environment before the inner runner starts.
-        runner.mState.enterTestClass(description);
-    }
-
-    /**
-     * Called after the inner runner finished.
-     */
-    public static void onAfterInnerRunnerFinished(
-            RavenwoodAwareTestRunner runner, Description description) throws Throwable {
-        Log.v(TAG, "onAfterInnerRunnerFinished: description=" + description);
-
-        RavenwoodTestStats.getInstance().onClassFinished(description);
-        runner.mState.exitTestClass();
-    }
-
-    /**
-     * Called before a test / class.
-     *
-     * Return false if it should be skipped.
-     */
-    public static boolean onBefore(RavenwoodAwareTestRunner runner, Description description,
-            Scope scope, Order order) throws Throwable {
-        Log.v(TAG, "onBefore: description=" + description + ", " + scope + ", " + order);
-
-        if (scope == Scope.Instance && order == Order.Outer) {
-            // Start of a test method.
-            runner.mState.enterTestMethod(description);
-        }
-
-        final var classDescription = runner.mState.getClassDescription();
-
-        // Class-level annotations are checked by the runner already, so we only check
-        // method-level annotations here.
-        if (scope == Scope.Instance && order == Order.Outer) {
-            if (!RavenwoodEnablementChecker.shouldEnableOnRavenwood(
-                    description, true)) {
-                RavenwoodTestStats.getInstance().onTestFinished(
-                        classDescription, description, Result.Skipped);
-                return false;
-            }
-        }
-        return true;
-    }
-
-    /**
-     * Called after a test / class.
-     *
-     * Return false if the exception should be ignored.
-     */
-    public static boolean onAfter(RavenwoodAwareTestRunner runner, Description description,
-            Scope scope, Order order, Throwable th) {
-        Log.v(TAG, "onAfter: description=" + description + ", " + scope + ", " + order + ", " + th);
-
-        final var classDescription = runner.mState.getClassDescription();
-
-        if (scope == Scope.Instance && order == Order.Outer) {
-            // End of a test method.
-            runner.mState.exitTestMethod();
-
-            final Result result;
-            if (th == null) {
-                result = Result.Passed;
-            } else if (th instanceof AssumptionViolatedException) {
-                result = Result.Skipped;
-            } else {
-                result = Result.Failed;
-            }
-
-            RavenwoodTestStats.getInstance().onTestFinished(classDescription, description, result);
-        }
-
-        // If RUN_DISABLED_TESTS is set, and the method did _not_ throw, make it an error.
-        if (RavenwoodRule.private$ravenwood().isRunningDisabledTests()
-                && scope == Scope.Instance && order == Order.Outer) {
-
-            boolean isTestEnabled = RavenwoodEnablementChecker.shouldEnableOnRavenwood(
-                    description, false);
-            if (th == null) {
-                // Test passed. Is the test method supposed to be enabled?
-                if (isTestEnabled) {
-                    // Enabled and didn't throw, okay.
-                    return true;
-                } else {
-                    // Disabled and didn't throw. We should report it.
-                    fail("Test wasn't included under Ravenwood, but it actually "
-                            + "passed under Ravenwood; consider updating annotations");
-                    return true; // unreachable.
-                }
-            } else {
-                // Test failed.
-                if (isTestEnabled) {
-                    // Enabled but failed. We should throw the exception.
-                    return true;
-                } else {
-                    // Disabled and failed. Expected. Don't throw.
-                    return false;
-                }
-            }
-        }
-        return true;
-    }
-
-    /**
-     * Called by {@link RavenwoodAwareTestRunner} to see if it should run a test class or not.
-     */
-    public static boolean shouldRunClassOnRavenwood(Class<?> clazz) {
-        return RavenwoodEnablementChecker.shouldRunClassOnRavenwood(clazz, true);
-    }
-
-    /**
-     * Called by RavenwoodRule.
-     */
-    public static void onRavenwoodRuleEnter(RavenwoodAwareTestRunner runner,
-            Description description, RavenwoodRule rule) throws Throwable {
-        Log.v(TAG, "onRavenwoodRuleEnter: description=" + description);
-
-        runner.mState.enterRavenwoodRule(rule);
-    }
-
-
-    /**
-     * Called by RavenwoodRule.
-     */
-    public static void onRavenwoodRuleExit(RavenwoodAwareTestRunner runner,
-            Description description, RavenwoodRule rule) throws Throwable {
-        Log.v(TAG, "onRavenwoodRuleExit: description=" + description);
-
-        runner.mState.exitRavenwoodRule(rule);
-    }
-}
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodConfigPrivate.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodConfigPrivate.java
new file mode 100644
index 000000000000..ffb642d60dcc
--- /dev/null
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodConfigPrivate.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright (C) 2024 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 android.platform.test.ravenwood;
+
+import android.annotation.Nullable;
+
+import java.util.function.BiConsumer;
+
+/**
+ * Contains Ravenwood private APIs.
+ */
+public class RavenwoodConfigPrivate {
+    private RavenwoodConfigPrivate() {
+    }
+
+    /**
+     * Set a listener for onCriticalError(), for testing. If a listener is set, we won't call
+     * System.exit().
+     */
+    public static void setCriticalErrorHandler(@Nullable BiConsumer<String, Throwable> handler) {
+        RavenwoodAwareTestRunner.sCriticalErrorHandler = handler;
+    }
+}
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunNotifier.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunNotifier.java
new file mode 100644
index 000000000000..69030350bf82
--- /dev/null
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunNotifier.java
@@ -0,0 +1,227 @@
+/*
+ * Copyright (C) 2024 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 android.platform.test.ravenwood;
+
+import android.util.Log;
+
+import org.junit.AssumptionViolatedException;
+import org.junit.runner.Description;
+import org.junit.runner.Result;
+import org.junit.runner.notification.Failure;
+import org.junit.runner.notification.RunListener;
+import org.junit.runner.notification.RunNotifier;
+import org.junit.runner.notification.StoppedByUserException;
+import org.junit.runners.model.MultipleFailureException;
+
+import java.util.ArrayList;
+import java.util.Stack;
+
+/**
+ * A run notifier that wraps another notifier and provides the following features:
+ * - Handle a failure that happened before testStarted and testEnded (typically that means
+ *   it's from @BeforeClass or @AfterClass, or a @ClassRule) and deliver it as if
+ *   individual tests in the class reported it. This is for b/364395552.
+ *
+ * - Logging.
+ */
+class RavenwoodRunNotifier extends RunNotifier {
+    private final RunNotifier mRealNotifier;
+
+    private final Stack<Description> mSuiteStack = new Stack<>();
+    private Description mCurrentSuite = null;
+    private final ArrayList<Throwable> mOutOfTestFailures = new ArrayList<>();
+
+    private boolean mBeforeTest = true;
+    private boolean mAfterTest = false;
+
+    RavenwoodRunNotifier(RunNotifier realNotifier) {
+        mRealNotifier = realNotifier;
+    }
+
+    private boolean isInTest() {
+        return !mBeforeTest && !mAfterTest;
+    }
+
+    @Override
+    public void addListener(RunListener listener) {
+        mRealNotifier.addListener(listener);
+    }
+
+    @Override
+    public void removeListener(RunListener listener) {
+        mRealNotifier.removeListener(listener);
+    }
+
+    @Override
+    public void addFirstListener(RunListener listener) {
+        mRealNotifier.addFirstListener(listener);
+    }
+
+    @Override
+    public void fireTestRunStarted(Description description) {
+        Log.i(RavenwoodAwareTestRunner.TAG, "testRunStarted: " + description);
+        mRealNotifier.fireTestRunStarted(description);
+    }
+
+    @Override
+    public void fireTestRunFinished(Result result) {
+        Log.i(RavenwoodAwareTestRunner.TAG, "testRunFinished: "
+                + result.getRunCount() + ","
+                + result.getFailureCount() + ","
+                + result.getAssumptionFailureCount() + ","
+                + result.getIgnoreCount());
+        mRealNotifier.fireTestRunFinished(result);
+    }
+
+    @Override
+    public void fireTestSuiteStarted(Description description) {
+        Log.i(RavenwoodAwareTestRunner.TAG, "testSuiteStarted: " + description);
+        mRealNotifier.fireTestSuiteStarted(description);
+
+        mBeforeTest = true;
+        mAfterTest = false;
+
+        // Keep track of the current suite, needed if the outer test is a Suite,
+        // in which case its children are test classes. (not test methods)
+        mCurrentSuite = description;
+        mSuiteStack.push(description);
+
+        mOutOfTestFailures.clear();
+    }
+
+    @Override
+    public void fireTestSuiteFinished(Description description) {
+        Log.i(RavenwoodAwareTestRunner.TAG, "testSuiteFinished: " + description);
+        mRealNotifier.fireTestSuiteFinished(description);
+
+        maybeHandleOutOfTestFailures();
+
+        mBeforeTest = true;
+        mAfterTest = false;
+
+        // Restore the upper suite.
+        mSuiteStack.pop();
+        mCurrentSuite = mSuiteStack.size() == 0 ? null : mSuiteStack.peek();
+    }
+
+    @Override
+    public void fireTestStarted(Description description) throws StoppedByUserException {
+        Log.i(RavenwoodAwareTestRunner.TAG, "testStarted: " + description);
+        mRealNotifier.fireTestStarted(description);
+
+        mAfterTest = false;
+        mBeforeTest = false;
+    }
+
+    @Override
+    public void fireTestFailure(Failure failure) {
+        Log.i(RavenwoodAwareTestRunner.TAG, "testFailure: " + failure);
+
+        if (isInTest()) {
+            mRealNotifier.fireTestFailure(failure);
+        } else {
+            mOutOfTestFailures.add(failure.getException());
+        }
+    }
+
+    @Override
+    public void fireTestAssumptionFailed(Failure failure) {
+        Log.i(RavenwoodAwareTestRunner.TAG, "testAssumptionFailed: " + failure);
+
+        if (isInTest()) {
+            mRealNotifier.fireTestAssumptionFailed(failure);
+        } else {
+            mOutOfTestFailures.add(failure.getException());
+        }
+    }
+
+    @Override
+    public void fireTestIgnored(Description description) {
+        Log.i(RavenwoodAwareTestRunner.TAG, "testIgnored: " + description);
+        mRealNotifier.fireTestIgnored(description);
+    }
+
+    @Override
+    public void fireTestFinished(Description description) {
+        Log.i(RavenwoodAwareTestRunner.TAG, "testFinished: " + description);
+        mRealNotifier.fireTestFinished(description);
+
+        mAfterTest = true;
+    }
+
+    @Override
+    public void pleaseStop() {
+        Log.w(RavenwoodAwareTestRunner.TAG, "pleaseStop:");
+        mRealNotifier.pleaseStop();
+    }
+
+    /**
+     * At the end of each Suite, we handle failures happened out of test methods.
+     * (typically in @BeforeClass or @AfterClasses)
+     *
+     * This is to work around b/364395552.
+     */
+    private boolean maybeHandleOutOfTestFailures() {
+        if (mOutOfTestFailures.size() == 0) {
+            return false;
+        }
+        Throwable th;
+        if (mOutOfTestFailures.size() == 1) {
+            th = mOutOfTestFailures.get(0);
+        } else {
+            th = new MultipleFailureException(mOutOfTestFailures);
+        }
+        if (mBeforeTest) {
+            reportBeforeTestFailure(mCurrentSuite, th);
+            return true;
+        }
+        if (mAfterTest) {
+            reportAfterTestFailure(th);
+            return true;
+        }
+        return false;
+    }
+
+    public void reportBeforeTestFailure(Description suiteDesc, Throwable th) {
+        // If a failure happens befere running any tests, we'll need to pretend
+        // as if each test in the suite reported the failure, to work around b/364395552.
+        for (var child : suiteDesc.getChildren()) {
+            if (child.isSuite()) {
+                // If the chiil is still a "parent" -- a test class or a test suite
+                // -- propagate to its children.
+                mRealNotifier.fireTestSuiteStarted(child);
+                reportBeforeTestFailure(child, th);
+                mRealNotifier.fireTestSuiteFinished(child);
+            } else {
+                mRealNotifier.fireTestStarted(child);
+                Failure f = new Failure(child, th);
+                if (th instanceof AssumptionViolatedException) {
+                    mRealNotifier.fireTestAssumptionFailed(f);
+                } else {
+                    mRealNotifier.fireTestFailure(f);
+                }
+                mRealNotifier.fireTestFinished(child);
+            }
+        }
+    }
+
+    public void reportAfterTestFailure(Throwable th) {
+        // Unfortunately, there's no good way to report it, so kill the own process.
+        RavenwoodAwareTestRunner.onCriticalError(
+                "Failures detected in @AfterClass, which would be swallowed by tradefed",
+                th);
+    }
+}
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunnerState.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
index 03513ab0a2af..ead4a849dcff 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
@@ -20,8 +20,8 @@ import static com.android.ravenwood.common.RavenwoodCommonUtils.ensureIsPublicMe
 import static org.junit.Assert.fail;
 
 import android.annotation.Nullable;
+import android.util.Log;
 
-import com.android.internal.annotations.GuardedBy;
 import com.android.ravenwood.common.RavenwoodRuntimeException;
 
 import org.junit.ClassRule;
@@ -29,9 +29,7 @@ import org.junit.Rule;
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 
-import java.io.IOException;
 import java.lang.reflect.Field;
-import java.util.WeakHashMap;
 
 /**
  * Used to store various states associated with the current test runner that's inly needed
@@ -45,10 +43,6 @@ import java.util.WeakHashMap;
 public final class RavenwoodRunnerState {
     private static final String TAG = "RavenwoodRunnerState";
 
-    @GuardedBy("sStates")
-    private static final WeakHashMap<RavenwoodAwareTestRunner, RavenwoodRunnerState> sStates =
-            new WeakHashMap<>();
-
     private final RavenwoodAwareTestRunner mRunner;
 
     /**
@@ -69,7 +63,8 @@ public final class RavenwoodRunnerState {
         return mClassDescription;
     }
 
-    public void enterTestClass(Description classDescription) throws IOException {
+    public void enterTestClass(Description classDescription) {
+        Log.i(TAG, "enterTestClass: description=" + classDescription);
         mClassDescription = classDescription;
 
         mHasRavenwoodRule = hasRavenwoodRule(mRunner.getTestClass().getJavaClass());
@@ -81,6 +76,7 @@ public final class RavenwoodRunnerState {
     }
 
     public void exitTestClass() {
+        Log.i(TAG, "exitTestClass: description=" + mClassDescription);
         if (mCurrentConfig != null) {
             try {
                 RavenwoodRuntimeEnvironmentController.reset();
@@ -98,7 +94,7 @@ public final class RavenwoodRunnerState {
         mMethodDescription = null;
     }
 
-    public void enterRavenwoodRule(RavenwoodRule rule) throws IOException {
+    public void enterRavenwoodRule(RavenwoodRule rule) {
         if (!mHasRavenwoodRule) {
             fail("If you have a RavenwoodRule in your test, make sure the field type is"
                     + " RavenwoodRule so Ravenwood can detect it.");
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
index c2806daf99a1..de4357c4e7c5 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
@@ -191,7 +191,7 @@ public class RavenwoodRuntimeEnvironmentController {
     /**
      * Initialize the environment.
      */
-    public static void init(RavenwoodConfig config) throws IOException {
+    public static void init(RavenwoodConfig config) {
         if (RAVENWOOD_VERBOSE_LOGGING) {
             Log.i(TAG, "init() called here", new RuntimeException("STACKTRACE"));
         }