1 files changed, 40 insertions, 7 deletions
diff --git a/test/junit/scala/tools/testing/ClearAfterClass.java b/test/junit/scala/tools/testing/ClearAfterClass.java
index 232d459c4e..95e170ec13 100644
--- a/test/junit/scala/tools/testing/ClearAfterClass.java
+++ b/test/junit/scala/tools/testing/ClearAfterClass.java
@@ -1,20 +1,53 @@
 package scala.tools.testing;
 
-import org.junit.AfterClass;
+import org.junit.ClassRule;
+import org.junit.rules.TestRule;
+import org.junit.runners.model.Statement;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
- * Extend this class to use JUnit's @AfterClass. This annotation only works on static methods,
+ * Extend this class to use JUnit's @ClassRule. This annotation only works on static methods,
  * which cannot be written in Scala.
  *
  * Example: {@link scala.tools.nsc.backend.jvm.opt.InlinerTest}
  */
 public class ClearAfterClass {
-    public static interface Clearable {
-        void clear();
+    private static Map<Class<?>, Map<String, Object>> cache = new ConcurrentHashMap<>();
+
+    @ClassRule
+    public static TestRule clearClassCache() {
+        return (statement, desc) -> new Statement() {
+            @Override
+            public void evaluate() throws Throwable {
+                ConcurrentHashMap<String, Object> perClassCache = new ConcurrentHashMap<>();
+                cache.put(desc.getTestClass(), perClassCache);
+                try {
+                    statement.evaluate();
+                } finally {
+                    perClassCache.values().forEach(ClearAfterClass::closeIfClosable);
+                    cache.remove(desc.getTestClass());
+                }
+            }
+        };
     }
 
-    public static Clearable stateToClear;
+    private static void closeIfClosable(Object o) {
+        if (o instanceof Closeable) {
+            try {
+                ((Closeable) o).close();
+            } catch (IOException e) {
+                // ignore
+            }
+        }
+    }
+
+    public <T> T cached(String key, scala.Function0<T> t) {
+        Map<String, Object> perClassCache = cache.get(getClass());
+        return (T) perClassCache.computeIfAbsent(key, s -> t.apply());
+    }
 
-    @AfterClass
-    public static void clearState() { stateToClear.clear(); }
 }