1 files changed, 141 insertions, 0 deletions
diff --git a/java/com/android/modules/utils/testing/StaticMockFixtureRule.java b/java/com/android/modules/utils/testing/StaticMockFixtureRule.java
new file mode 100644
index 0000000..0cdc9a2
--- /dev/null
+++ b/java/com/android/modules/utils/testing/StaticMockFixtureRule.java
@@ -0,0 +1,141 @@
+/*
+ * Copyright (C) 2019 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 com.android.modules.utils.testing;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+
+import com.android.dx.mockito.inline.extended.StaticMockitoSession;
+import com.android.dx.mockito.inline.extended.StaticMockitoSessionBuilder;
+
+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.quality.Strictness;
+
+import java.util.function.Supplier;
+
+/**
+ * <p>StaticMockFixtureRule is a {@link TestRule} that wraps one or more {@link StaticMockFixture}s
+ * to set them up and tear it down automatically. This works well when you have no other static
+ * mocks than the ones supported by their respective {@link StaticMockFixture}s.</p>
+ *
+ * <p>StaticMockFixtureRule should be defined as a rule on your test so it can clean up after
+ * itself. Like the following:</p>
+ * <pre class="prettyprint">
+*  public final StaticMockFixture mStaticMockFixtures = ...;
+ * &#064;Rule
+ * public final StaticMockFixtureRule mStaticMockFixtureRule =
+ *     new StaticMockFixtureRule(mStaticMockFixtures);
+ * </pre>
+ */
+public class StaticMockFixtureRule implements TestRule {
+    private StaticMockitoSession mMockitoSession;
+    private StaticMockFixture[] mStaticMockFixtures;
+    private Supplier<? extends StaticMockFixture>[] mSupplier;
+
+    /**
+     * Constructs a StaticMockFixtureRule that always uses the same {@link StaticMockFixture}
+     * instance(s).
+     *
+     * @param staticMockFixtures the {@link StaticMockFixture}(s) to use.
+     */
+    public StaticMockFixtureRule(StaticMockFixture... staticMockFixtures) {
+        mStaticMockFixtures = staticMockFixtures;
+        mSupplier = null;
+    }
+
+    /**
+     * Constructs a StaticMockFixtureRule that retrieves a new {@link StaticMockFixture} instance
+     * from one or more {@link Supplier<? extends   StaticMockFixture  >}s for each test invocation.
+     *
+     * @param supplier the {@link Supplier<? extends   StaticMockFixture  >}(s) that will supply the
+     * {@link StaticMockFixture}(s).
+     */
+    @SafeVarargs
+    public StaticMockFixtureRule(Supplier<? extends StaticMockFixture>... supplier) {
+        mStaticMockFixtures = null;
+        mSupplier = supplier;
+    }
+
+    @Override
+    public Statement apply(Statement base, Description description) {
+        StaticMockitoSessionBuilder sessionBuilder = getSessionBuilder();
+
+        if (mSupplier != null) {
+            mStaticMockFixtures = new StaticMockFixture[mSupplier.length];
+            for (int i = 0; i < mSupplier.length; i++) {
+                mStaticMockFixtures[i] = mSupplier[i].get();
+            }
+        }
+
+        for (int i = 0; i < mStaticMockFixtures.length; i++) {
+            sessionBuilder = mStaticMockFixtures[i].setUpMockedClasses(sessionBuilder);
+        }
+
+        mMockitoSession = sessionBuilder.startMocking();
+
+        for (int i = 0; i < mStaticMockFixtures.length; i++) {
+            mStaticMockFixtures[i].setUpMockBehaviors();
+        }
+
+        return new TestWatcher() {
+            @Override
+            protected void succeeded(Description description) {
+                tearDown(null);
+            }
+
+            @Override
+            protected void skipped(AssumptionViolatedException e, Description description) {
+                tearDown(e);
+            }
+
+            @Override
+            protected void failed(Throwable e, Description description) {
+                tearDown(e);
+            }
+        }.apply(base, description);
+    }
+
+    /**
+     * This allows overriding the creation of the builder for a new {@link StaticMockitoSession}.
+     * Mainly for testing, but also useful if you have other requirements for the session.
+     *
+     * @return a new {@link StaticMockitoSessionBuilder}.
+     */
+    public StaticMockitoSessionBuilder getSessionBuilder() {
+        return mockitoSession().strictness(Strictness.LENIENT);
+    }
+
+    private void tearDown(Throwable e) {
+        mMockitoSession.finishMocking(e);
+
+        for (int i = mStaticMockFixtures.length - 1; i >= 0; i--) {
+            mStaticMockFixtures[i].tearDown();
+            if (mSupplier != null) {
+                mStaticMockFixtures[i] = null;
+            }
+        }
+
+        if (mSupplier != null) {
+            mStaticMockFixtures = null;
+        }
+
+        mMockitoSession = null;
+    }
+}