Check that the required state has been initialized before running a test

Author: Kiskae

src/main/java/com/google/testing/compile/CompilationExtension.java

@@ -24,6 +24,7 @@
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import org.junit.jupiter.api.extension.AfterAllCallback;
 import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.Extension;
 import org.junit.jupiter.api.extension.ExtensionContext;
@@ -67,7 +68,7 @@
  * @author David van Leusen
  */
 public class CompilationExtension
-    implements BeforeAllCallback, AfterAllCallback, ParameterResolver {
+    implements BeforeAllCallback, BeforeEachCallback, AfterAllCallback, ParameterResolver {
   private static final JavaFileObject DUMMY =
       JavaFileObjects.forSourceLines("Dummy", "final class Dummy {}");
   private static final ExtensionContext.Namespace NAMESPACE =
@@ -130,6 +131,15 @@ protected boolean onAdvance(int phase, int parties) {
     }
   }
 
+  @Override
+  public void beforeEach(ExtensionContext extensionContext) {
+    checkState(
+        PHASER_KEY.get(extensionContext.getStore(NAMESPACE)) != null,
+        "CompilationExtension is only available as a class-level extension. " +
+            "Using it as an instance-level extension through @RegisterExtension is not supported"
+    );
+  }
+
   @Override
   public void afterAll(ExtensionContext context) throws Exception {
     final ExtensionContext.Store store = context.getStore(NAMESPACE);