Fixed AbstractExtendedMockitoRule when test throws AVE. am: 6bb0293fbf am: 58e31437ae

Original change: https://android-review.googlesource.com/c/platform/frameworks/libs/modules-utils/+/2719505

Change-Id: I804d7392f0ba6995218f4277c48fd5b175599ebd
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

1 files changed, 30 insertions, 18 deletions

diff --git a/java/com/android/modules/utils/testing/AbstractExtendedMockitoRule.java b/java/com/android/modules/utils/testing/AbstractExtendedMockitoRule.java
index 05e2894..c63a8c5 100644
--- a/java/com/android/modules/utils/testing/AbstractExtendedMockitoRule.java
+++ b/java/com/android/modules/utils/testing/AbstractExtendedMockitoRule.java
@@ -24,9 +24,7 @@ import com.android.dx.mockito.inline.extended.StaticMockitoSessionBuilder;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.modules.utils.testing.AbstractExtendedMockitoRule.AbstractBuilder;
 
-import org.junit.AssumptionViolatedException;
 import org.junit.rules.TestRule;
-import org.junit.rules.TestWatcher;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
 import org.mockito.Mockito;
@@ -88,27 +86,41 @@ public abstract class AbstractExtendedMockitoRule<R extends AbstractExtendedMock
 
     @Override
     public Statement apply(Statement base, Description description) {
-        createMockitoSession(description);
-
-        return new TestWatcher() {
-            @Override
-            protected void succeeded(Description description) {
-                tearDown(description, /* e=*/ null);
-            }
-
-            @Override
-            protected void skipped(AssumptionViolatedException e, Description description) {
-                tearDown(description, e);
-            }
-
+        return new Statement() {
             @Override
-            protected void failed(Throwable e, Description description) {
-                tearDown(description, e);
+            public void evaluate() throws Throwable {
+                createMockitoSession(description);
+                Throwable error = null;
+                try {
+                    // TODO(b/296937563): need to add unit tests that make sure the session is
+                    // always closed
+                    base.evaluate();
+                } catch (Throwable t) {
+                    error = t;
+                }
+                try {
+                    tearDown(description, error);
+                } catch (Throwable t) {
+                    if (error != null) {
+                        Log.e(TAG, "Teardown failed for " + description.getDisplayName()
+                            + ", but not throwing it because test also threw (" + error + ")", t);
+                    } else {
+                        error = t;
+                    }
+                }
+                if (error != null) {
+                    // TODO(b/296937563): ideally should also add unit tests to make sure the
+                    // test error is thrown (in case tearDown() above fails)
+                    throw error;
+                }
             }
-        }.apply(base, description);
+        };
     }
 
     private void createMockitoSession(Description description) {
+        // TODO(b/296937563): might be prudent to save the session statically so it's explicitly
+        // closed in case it fails to be created again if for some reason it was not closed by us
+        // (although that should not happen)
         Log.v(TAG, "Creating session builder with strictness " + mStrictness);
         StaticMockitoSessionBuilder mSessionBuilder = mockitoSession().strictness(mStrictness);