[Ravenwood] Update RATR to setup environment ASAP

- Remove the workaround to record constructor exceptions and throw
  later, runner errors are handled properly after aosp/3310766
- Initialize Ravenwood's environment as soon as the real inner runner is
  instantiated, as in some cases getDescription() itself needs env setup
- Make the entire environment tied to a RATR instance
- Add new tests to make sure the early environment setup is working

Flag: EXEMPT host test change only
Bug: 356918135
Test: $ANDROID_BUILD_TOP/frameworks/base/ravenwood/scripts/run-ravenwood-tests.sh
Change-Id: I301607fd6602649172f16991d6d5ae1b577c47cd

3 files changed, 111 insertions, 138 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
index 30a653d2da76..66a6890a23b0 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java
@@ -38,7 +38,6 @@ 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;
@@ -94,7 +93,7 @@ public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase
     /** Keeps track of the runner on the current thread. */
     private static final ThreadLocal<RavenwoodAwareTestRunner> sCurrentRunner = new ThreadLocal<>();
 
-    static RavenwoodAwareTestRunner getCurrentRunner() {
+    private static RavenwoodAwareTestRunner getCurrentRunner() {
         var runner = sCurrentRunner.get();
         if (runner == null) {
             throw new RuntimeException("Current test runner not set!");
@@ -102,11 +101,9 @@ public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase
         return runner;
     }
 
-    private final Class<?> mTestJavaClass;
+    final Class<?> mTestJavaClass;
+    private final Runner mRealRunner;
     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
@@ -114,50 +111,37 @@ public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase
      */
     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);
+        /*
+         * 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);
+            return;
+        }
+
+        mTestClass = new TestClass(testClass);
+
+        Log.v(TAG, "RavenwoodAwareTestRunner starting for " + testClass.getCanonicalName());
 
-            Log.v(TAG, "RavenwoodAwareTestRunner starting for " + testClass.getCanonicalName());
+        // This is needed to make AndroidJUnit4ClassRunner happy.
+        InstrumentationRegistry.registerInstance(null, Bundle.EMPTY);
 
-            onRunnerInitializing();
+        // Hook point to allow more customization.
+        runAnnotatedMethodsOnRavenwood(RavenwoodTestRunnerInitializing.class, null);
 
-            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");
+        mRealRunner = instantiateRealRunner(mTestClass);
 
-            // This is for testing if tradefed is fixed.
-            if ("1".equals(System.getenv("RAVENWOOD_THROW_EXCEPTION_IN_TEST_RUNNER"))) {
-                throw th;
-            }
-        }
+        mState.enterTestRunner();
     }
 
     @Override
@@ -165,23 +149,11 @@ public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase
         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)) {
+        for (var method : mTestClass.getAnnotatedMethods(annotationClass)) {
             ensureIsPublicVoidMethod(method.getMethod(), /* isStatic=*/ instance == null);
 
             var methodDesc = method.getDeclaringClass().getName() + "."
@@ -195,11 +167,6 @@ public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase
     }
 
     @Override
-    public Description getDescription() {
-        return mDescription;
-    }
-
-    @Override
     public void run(RunNotifier realNotifier) {
         final var notifier = new RavenwoodRunNotifier(realNotifier);
         final var description = getDescription();
@@ -216,10 +183,6 @@ public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase
             dumpDescription(description);
         }
 
-        if (maybeReportExceptionFromConstructor(notifier)) {
-            return;
-        }
-
         // TODO(b/365976974): handle nested classes better
         final boolean skipRunnerHook =
                 mRealRunnerTakesRunnerBuilder && mRealRunner instanceof Suite;
@@ -228,7 +191,7 @@ public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase
         try {
             if (!skipRunnerHook) {
                 try {
-                    mState.enterTestClass(description);
+                    mState.enterTestClass();
                 } catch (Throwable th) {
                     notifier.reportBeforeTestFailure(description, th);
                     return;
@@ -251,18 +214,6 @@ public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase
         }
     }
 
-    /** 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() {
@@ -338,7 +289,7 @@ public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase
             mState.enterTestMethod(description);
         }
 
-        final var classDescription = mState.getClassDescription();
+        final var classDescription = getDescription();
 
         // Class-level annotations are checked by the runner already, so we only check
         // method-level annotations here.
@@ -360,7 +311,7 @@ public final class RavenwoodAwareTestRunner extends RavenwoodAwareTestRunnerBase
     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();
+        final var classDescription = getDescription();
 
         if (scope == Scope.Instance && order == Order.Outer) {
             // End of a test method.
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 ead4a849dcff..ec00e8fea887 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
@@ -52,37 +52,65 @@ public final class RavenwoodRunnerState {
         mRunner = runner;
     }
 
-    private Description mClassDescription;
-    private Description mMethodDescription;
-
+    /**
+     * The RavenwoodConfig used to configure the current Ravenwood environment.
+     * This can either come from mConfig or mRule.
+     */
     private RavenwoodConfig mCurrentConfig;
-    private RavenwoodRule mCurrentRule;
+    /**
+     * The RavenwoodConfig declared in the test class
+     */
+    private RavenwoodConfig mConfig;
+    /**
+     * The RavenwoodRule currently in effect, declared in the test class
+     */
+    private RavenwoodRule mRule;
     private boolean mHasRavenwoodRule;
+    private Description mMethodDescription;
 
-    public Description getClassDescription() {
-        return mClassDescription;
+    public RavenwoodConfig getConfig() {
+        return mCurrentConfig;
     }
 
-    public void enterTestClass(Description classDescription) {
-        Log.i(TAG, "enterTestClass: description=" + classDescription);
-        mClassDescription = classDescription;
+    public void enterTestRunner() {
+        Log.i(TAG, "enterTestRunner: " + mRunner);
+
+        mHasRavenwoodRule = hasRavenwoodRule(mRunner.mTestJavaClass);
+        mConfig = extractConfiguration(mRunner.mTestJavaClass);
 
-        mHasRavenwoodRule = hasRavenwoodRule(mRunner.getTestClass().getJavaClass());
-        mCurrentConfig = extractConfiguration(mRunner.getTestClass().getJavaClass());
+        if (mConfig != null) {
+            if (mHasRavenwoodRule) {
+                fail("RavenwoodConfig and RavenwoodRule cannot be used in the same class."
+                        + " Suggest migrating to RavenwoodConfig.");
+            }
+            mCurrentConfig = mConfig;
+        } else if (!mHasRavenwoodRule) {
+            // If no RavenwoodConfig and no RavenwoodRule, use a default config
+            mCurrentConfig = new RavenwoodConfig.Builder().build();
+        }
 
         if (mCurrentConfig != null) {
-            RavenwoodRuntimeEnvironmentController.init(mCurrentConfig);
+            RavenwoodRuntimeEnvironmentController.init(mRunner);
         }
     }
 
-    public void exitTestClass() {
-        Log.i(TAG, "exitTestClass: description=" + mClassDescription);
+    public void enterTestClass() {
+        Log.i(TAG, "enterTestClass: " + mRunner.mTestJavaClass.getName());
+
         if (mCurrentConfig != null) {
-            try {
+            RavenwoodRuntimeEnvironmentController.init(mRunner);
+        }
+    }
+
+    public void exitTestClass() {
+        Log.i(TAG, "exitTestClass: " + mRunner.mTestJavaClass.getName());
+        try {
+            if (mCurrentConfig != null) {
                 RavenwoodRuntimeEnvironmentController.reset();
-            } finally {
-                mClassDescription = null;
             }
+        } finally {
+            mConfig = null;
+            mRule = null;
         }
     }
 
@@ -92,6 +120,7 @@ public final class RavenwoodRunnerState {
 
     public void exitTestMethod() {
         mMethodDescription = null;
+        RavenwoodRuntimeEnvironmentController.reinit();
     }
 
     public void enterRavenwoodRule(RavenwoodRule rule) {
@@ -99,48 +128,32 @@ public final class RavenwoodRunnerState {
             fail("If you have a RavenwoodRule in your test, make sure the field type is"
                     + " RavenwoodRule so Ravenwood can detect it.");
         }
-        if (mCurrentConfig != null) {
-            fail("RavenwoodConfig and RavenwoodRule cannot be used in the same class."
-                    + " Suggest migrating to RavenwoodConfig.");
-        }
-        if (mCurrentRule != null) {
+        if (mRule != null) {
             fail("Multiple nesting RavenwoodRule's are detected in the same class,"
                     + " which is not supported.");
         }
-        mCurrentRule = rule;
-        RavenwoodRuntimeEnvironmentController.init(rule.getConfiguration());
+        mRule = rule;
+        if (mCurrentConfig == null) {
+            mCurrentConfig = rule.getConfiguration();
+        }
+        RavenwoodRuntimeEnvironmentController.init(mRunner);
     }
 
     public void exitRavenwoodRule(RavenwoodRule rule) {
-        if (mCurrentRule != rule) {
-            return; // This happens if the rule did _not_ take effect somehow.
-        }
-
-        try {
-            RavenwoodRuntimeEnvironmentController.reset();
-        } finally {
-            mCurrentRule = null;
+        if (mRule != rule) {
+            fail("RavenwoodRule did not take effect.");
         }
+        mRule = null;
     }
 
     /**
      * @return a configuration from a test class, if any.
      */
     @Nullable
-    private RavenwoodConfig extractConfiguration(Class<?> testClass) {
+    private static RavenwoodConfig extractConfiguration(Class<?> testClass) {
         var field = findConfigurationField(testClass);
         if (field == null) {
-            if (mHasRavenwoodRule) {
-                // Should be handled by RavenwoodRule
-                return null;
-            }
-
-            // If no RavenwoodConfig and no RavenwoodRule, return a default config
-            return new RavenwoodConfig.Builder().build();
-        }
-        if (mHasRavenwoodRule) {
-            fail("RavenwoodConfig and RavenwoodRule cannot be used in the same class."
-                    + " Suggest migrating to RavenwoodConfig.");
+            return null;
         }
 
         try {
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 de4357c4e7c5..9002e40bba32 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
@@ -137,7 +137,7 @@ public class RavenwoodRuntimeEnvironmentController {
         return res;
     }
 
-    private static RavenwoodConfig sConfig;
+    private static RavenwoodAwareTestRunner sRunner;
     private static RavenwoodSystemProperties sProps;
     private static boolean sInitialized = false;
 
@@ -171,6 +171,10 @@ public class RavenwoodRuntimeEnvironmentController {
         // Redirect stdout/stdin to liblog.
         RuntimeInit.redirectLogStreams();
 
+        // Touch some references early to ensure they're <clinit>'ed
+        Objects.requireNonNull(Build.TYPE);
+        Objects.requireNonNull(Build.VERSION.SDK);
+
         if (RAVENWOOD_VERBOSE_LOGGING) {
             RavenwoodCommonUtils.log(TAG, "Force enabling verbose logging");
             try {
@@ -191,12 +195,19 @@ public class RavenwoodRuntimeEnvironmentController {
     /**
      * Initialize the environment.
      */
-    public static void init(RavenwoodConfig config) {
+    public static void init(RavenwoodAwareTestRunner runner) {
         if (RAVENWOOD_VERBOSE_LOGGING) {
-            Log.i(TAG, "init() called here", new RuntimeException("STACKTRACE"));
+            Log.i(TAG, "init() called here: " + runner, new RuntimeException("STACKTRACE"));
         }
+        if (sRunner == runner) {
+            return;
+        }
+        if (sRunner != null) {
+            reset();
+        }
+        sRunner = runner;
         try {
-            initInner(config);
+            initInner(runner.mState.getConfig());
         } catch (Exception th) {
             Log.e(TAG, "init() failed", th);
             reset();
@@ -205,10 +216,6 @@ public class RavenwoodRuntimeEnvironmentController {
     }
 
     private static void initInner(RavenwoodConfig config) throws IOException {
-        if (sConfig != null) {
-            throw new RavenwoodRuntimeException("Internal error: init() called without reset()");
-        }
-        sConfig = config;
         if (ENABLE_UNCAUGHT_EXCEPTION_DETECTION) {
             maybeThrowPendingUncaughtException(false);
             Thread.setDefaultUncaughtExceptionHandler(sUncaughtExceptionHandler);
@@ -216,7 +223,7 @@ public class RavenwoodRuntimeEnvironmentController {
 
         android.os.Process.init$ravenwood(config.mUid, config.mPid);
         sOriginalIdentityToken = Binder.clearCallingIdentity();
-        Binder.restoreCallingIdentity(packBinderIdentityToken(false, config.mUid, config.mPid));
+        reinit();
         setSystemProperties(config.mSystemProperties);
 
         ServiceManager.init$ravenwood();
@@ -269,9 +276,7 @@ public class RavenwoodRuntimeEnvironmentController {
         config.mInstContext = instContext;
         config.mTargetContext = targetContext;
 
-        final Supplier<Resources> systemResourcesLoader = () -> {
-            return config.mState.loadResources(null);
-        };
+        final Supplier<Resources> systemResourcesLoader = () -> config.mState.loadResources(null);
 
         config.mState.mSystemServerContext =
                 new RavenwoodContext(ANDROID_PACKAGE_NAME, main, systemResourcesLoader);
@@ -288,10 +293,14 @@ public class RavenwoodRuntimeEnvironmentController {
                     RavenwoodRuntimeEnvironmentController::dumpStacks,
                     TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
         }
+    }
 
-        // Touch some references early to ensure they're <clinit>'ed
-        Objects.requireNonNull(Build.TYPE);
-        Objects.requireNonNull(Build.VERSION.SDK);
+    /**
+     * Partially re-initialize after each test method invocation
+     */
+    public static void reinit() {
+        var config = sRunner.mState.getConfig();
+        Binder.restoreCallingIdentity(packBinderIdentityToken(false, config.mUid, config.mPid));
     }
 
     /**
@@ -301,11 +310,11 @@ public class RavenwoodRuntimeEnvironmentController {
         if (RAVENWOOD_VERBOSE_LOGGING) {
             Log.i(TAG, "reset() called here", new RuntimeException("STACKTRACE"));
         }
-        if (sConfig == null) {
+        if (sRunner == null) {
             throw new RavenwoodRuntimeException("Internal error: reset() already called");
         }
-        var config = sConfig;
-        sConfig = null;
+        var config = sRunner.mState.getConfig();
+        sRunner = null;
 
         if (ENABLE_TIMEOUT_STACKS) {
             sPendingTimeout.cancel(false);