Make entry-based extensions execute around classes

src/main/java/org/junitpioneer/jupiter/AbstractEntryBasedExtension.java

@@ -15,6 +15,7 @@
 
 import java.lang.annotation.Annotation;
 import java.lang.reflect.AnnotatedElement;
+import java.lang.reflect.Method;
 import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.Type;
 import java.util.Collection;
@@ -25,9 +26,12 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.stream.Stream;
 
+import org.junit.jupiter.api.extension.AfterAllCallback;
 import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionConfigurationException;
 import org.junit.jupiter.api.extension.ExtensionContext;
@@ -46,7 +50,14 @@
  * @param <S> The set annotation type.
  */
 abstract class AbstractEntryBasedExtension<K, V, C extends Annotation, S extends Annotation>
-		implements BeforeEachCallback, AfterEachCallback {
+		implements BeforeEachCallback, AfterEachCallback, BeforeAllCallback, AfterAllCallback {
+
+	@Override
+	public void beforeAll(ExtensionContext context) {
+		List<ExtensionContext> contexts = PioneerUtils.findAllContexts(context);
+		Collections.reverse(contexts);
+		contexts.forEach(c -> clearAndSetEntries(c, ApplyMode.CLASS));
+	}
 
 	@Override
 	public void beforeEach(ExtensionContext context) {
@@ -57,17 +68,21 @@ public void beforeEach(ExtensionContext context) {
 		 */
 		List<ExtensionContext> contexts = PioneerUtils.findAllContexts(context);
 		Collections.reverse(contexts);
-		contexts.forEach(this::clearAndSetEntries);
+		contexts.forEach(c -> clearAndSetEntries(c, ApplyMode.TEST));
 	}
 
-	private void clearAndSetEntries(ExtensionContext context) {
+	private void clearAndSetEntries(ExtensionContext context, ApplyMode mode) {
+		if (context.getTestMethod().isPresent() && classApplyModeAnnotationsOnMethod(context.getRequiredTestMethod())) {
+			throw new ExtensionConfigurationException("Test-level annotations can not have ApplyMode.CLASS");
+		}
+
 		context.getElement().ifPresent(element -> {
 			Set<K> entriesToClear;
 			Map<K, V> entriesToSet;
 
 			try {
-				entriesToClear = findEntriesToClear(element);
-				entriesToSet = findEntriesToSet(element);
+				entriesToClear = findEntriesToClear(element, mode);
+				entriesToSet = findEntriesToSet(element, mode);
 				preventClearAndSetSameEntries(entriesToClear, entriesToSet.keySet());
 			}
 			catch (IllegalStateException ex) {
@@ -78,20 +93,28 @@ private void clearAndSetEntries(ExtensionContext context) {
 				return;
 
 			reportWarning(context);
-			storeOriginalEntries(context, entriesToClear, entriesToSet.keySet());
+			storeOriginalEntries(context, entriesToClear, entriesToSet.keySet(), mode);
 			clearEntries(entriesToClear);
 			setEntries(entriesToSet);
 		});
 	}
 
-	private Set<K> findEntriesToClear(AnnotatedElement element) {
+	private boolean classApplyModeAnnotationsOnMethod(Method method) {
+		return !(findEntriesToClear(method, ApplyMode.CLASS).isEmpty()
+				&& findEntriesToSet(method, ApplyMode.CLASS).isEmpty());
+	}
+
+	private Set<K> findEntriesToClear(AnnotatedElement element, ApplyMode mode) {
 		return findAnnotations(element, getClearAnnotationType())
+				.filter(filterClearAnnotationsByMode(mode))
 				.map(clearKeyMapper())
 				.collect(PioneerUtils.distinctToSet());
 	}
 
-	private Map<K, V> findEntriesToSet(AnnotatedElement element) {
-		return findAnnotations(element, getSetAnnotationType()).collect(toMap(setKeyMapper(), setValueMapper()));
+	private Map<K, V> findEntriesToSet(AnnotatedElement element, ApplyMode mode) {
+		return findAnnotations(element, getSetAnnotationType())
+				.filter(filterSetAnnotationsByMode(mode))
+				.collect(toMap(setKeyMapper(), setValueMapper()));
 	}
 
 	private <A extends Annotation> Stream<A> findAnnotations(AnnotatedElement element, Class<A> clazz) {
@@ -125,8 +148,8 @@ private void preventClearAndSetSameEntries(Collection<K> entriesToClear, Collect
 	}
 
 	private void storeOriginalEntries(ExtensionContext context, Collection<K> entriesToClear,
-			Collection<K> entriesToSet) {
-		getStore(context).put(getStoreKey(context), new EntriesBackup(entriesToClear, entriesToSet));
+			Collection<K> entriesToSet, ApplyMode mode) {
+		getStore(context).put(getStoreKey(context, mode), new EntriesBackup(entriesToClear, entriesToSet));
 	}
 
 	private void clearEntries(Collection<K> entriesToClear) {
@@ -138,21 +161,28 @@ private void setEntries(Map<K, V> entriesToSet) {
 	}
 
 	@Override
-	public void afterEach(ExtensionContext context) throws Exception {
+	public void afterEach(ExtensionContext context) {
 		// apply from innermost to outermost
-		PioneerUtils.findAllContexts(context).forEach(this::restoreOriginalEntries);
+		PioneerUtils.findAllContexts(context).forEach(c -> restoreOriginalEntries(c, ApplyMode.TEST));
 	}
 
-	private void restoreOriginalEntries(ExtensionContext context) {
-		getStore(context).getOrDefault(getStoreKey(context), EntriesBackup.class, new EntriesBackup()).restoreBackup();
+	@Override
+	public void afterAll(ExtensionContext context) {
+		PioneerUtils.findAllContexts(context).forEach(c -> restoreOriginalEntries(c, ApplyMode.CLASS));
+	}
+
+	private void restoreOriginalEntries(ExtensionContext context, ApplyMode mode) {
+		getStore(context)
+				.getOrDefault(getStoreKey(context, mode), EntriesBackup.class, new EntriesBackup())
+				.restoreBackup();
 	}
 
 	private Store getStore(ExtensionContext context) {
 		return context.getStore(Namespace.create(getClass()));
 	}
 
-	private Object getStoreKey(ExtensionContext context) {
-		return context.getUniqueId();
+	private Object getStoreKey(ExtensionContext context, ApplyMode mode) {
+		return context.getUniqueId() + mode.name();
 	}
 
 	private class EntriesBackup {
@@ -181,6 +211,16 @@ public void restoreBackup() {
 
 	}
 
+	/**
+	 * @return Filter function based on the actual {@link ApplyMode} for clear annotations.
+	 */
+	protected abstract Predicate<C> filterClearAnnotationsByMode(ApplyMode mode);
+
+	/**
+	 * @return Filter function based on the actual {@link ApplyMode} for clear annotations.
+	 */
+	protected abstract Predicate<S> filterSetAnnotationsByMode(ApplyMode mode);
+
 	/**
 	 * @return Mapper function to get the key from a clear annotation.
 	 */

src/main/java/org/junitpioneer/jupiter/ApplyMode.java

@@ -0,0 +1,15 @@
+/*
+ * Copyright 2016-2022 the original author or authors.
+ *
+ * All rights reserved. This program and the accompanying materials are
+ * made available under the terms of the Eclipse Public License v2.0 which
+ * accompanies this distribution and is available at
+ *
+ * http://www.eclipse.org/legal/epl-v20.html
+ */
+
+package org.junitpioneer.jupiter;
+
+public enum ApplyMode {
+	TEST, CLASS
+}

src/main/java/org/junitpioneer/jupiter/ClearEnvironmentVariable.java

@@ -64,6 +64,11 @@
 	 */
 	String key();
 
+	/**
+	 * Optional configuration option when to apply the extension.
+	 */
+	ApplyMode mode() default ApplyMode.TEST;
+
 	/**
 	 * Containing annotation of repeatable {@code @ClearEnvironmentVariable}.
 	 */

src/main/java/org/junitpioneer/jupiter/ClearSystemProperty.java

@@ -56,6 +56,11 @@
 	 */
 	String key();
 
+	/**
+	 * Optional configuration option when to apply the extension.
+	 */
+	ApplyMode mode() default ApplyMode.TEST;
+
 	/**
 	 * Containing annotation of repeatable {@code @ClearSystemProperty}.
 	 */

src/main/java/org/junitpioneer/jupiter/EnvironmentVariableExtension.java

@@ -12,6 +12,7 @@
 
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
+import java.util.function.Predicate;
 
 import org.junit.jupiter.api.extension.ExtensionContext;
 
@@ -23,6 +24,16 @@ class EnvironmentVariableExtension
 	static final String WARNING_KEY = EnvironmentVariableExtension.class.getSimpleName();
 	static final String WARNING_VALUE = "This extension uses reflection to mutate JDK-internal state, which is fragile. Check the Javadoc or documentation for more details.";
 
+	@Override
+	protected Predicate<ClearEnvironmentVariable> filterClearAnnotationsByMode(ApplyMode mode) {
+		return annotation -> annotation.mode().equals(mode);
+	}
+
+	@Override
+	protected Predicate<SetEnvironmentVariable> filterSetAnnotationsByMode(ApplyMode mode) {
+		return annotation -> annotation.mode().equals(mode);
+	}
+
 	@Override
 	protected Function<ClearEnvironmentVariable, String> clearKeyMapper() {
 		return ClearEnvironmentVariable::key;

src/main/java/org/junitpioneer/jupiter/SetEnvironmentVariable.java

@@ -71,6 +71,11 @@
 	 */
 	String value();
 
+	/**
+	 * Optional configuration option when to apply the extension.
+	 */
+	ApplyMode mode() default ApplyMode.TEST;
+
 	/**
 	 * Containing annotation of repeatable {@code @SetEnvironmentVariable}.
 	 */

src/main/java/org/junitpioneer/jupiter/SetSystemProperty.java

@@ -63,6 +63,11 @@
 	 */
 	String value();
 
+	/**
+	 * Optional configuration option when to apply the extension.
+	 */
+	ApplyMode mode() default ApplyMode.TEST;
+
 	/**
 	 * Containing annotation of repeatable {@code @SetSystemProperty}.
 	 */

src/main/java/org/junitpioneer/jupiter/SystemPropertyExtension.java

@@ -11,10 +11,21 @@
 package org.junitpioneer.jupiter;
 
 import java.util.function.Function;
+import java.util.function.Predicate;
 
 class SystemPropertyExtension
 		extends AbstractEntryBasedExtension<String, String, ClearSystemProperty, SetSystemProperty> {
 
+	@Override
+	protected Predicate<ClearSystemProperty> filterClearAnnotationsByMode(ApplyMode mode) {
+		return annotation -> annotation.mode().equals(mode);
+	}
+
+	@Override
+	protected Predicate<SetSystemProperty> filterSetAnnotationsByMode(ApplyMode mode) {
+		return annotation -> annotation.mode().equals(mode);
+	}
+
 	@Override
 	protected Function<ClearSystemProperty, String> clearKeyMapper() {
 		return ClearSystemProperty::key;

src/test/java/org/junitpioneer/jupiter/EnvironmentVariableExtensionTests.java

@@ -247,7 +247,13 @@ class ResettingEnvironmentVariableTests {
 		@Nested
 		@SetEnvironmentVariable(key = "set envvar A", value = "newer A")
 		@TestInstance(TestInstance.Lifecycle.PER_CLASS)
-		class ResettingEnvironmentVariableNestedTests {
+		class ResettingEnvironmentVariableAfterEachNestedTests {
+
+			@BeforeAll
+			@ReadsEnvironmentVariable
+			void changeShouldNotBeVisible() {
+				assertThat(System.getenv("set envvar A")).isEqualTo("old A");
+			}
 
 			@Test
 			@SetEnvironmentVariable(key = "set envvar A", value = "newest A")
@@ -263,6 +269,31 @@ void resetAfterTestMethodExecution() {
 
 		}
 
+		@Nested
+		@SetEnvironmentVariable(key = "set envvar A", value = "newer A", mode = ApplyMode.CLASS)
+		@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+		class ResettingEnvironmentVariableAfterAllNestedTests {
+
+			@BeforeAll
+			@ReadsEnvironmentVariable
+			void changeShouldBeVisible() {
+				assertThat(System.getenv("set envvar A")).isEqualTo("newer A");
+			}
+
+			@Test
+			@SetEnvironmentVariable(key = "set envvar A", value = "newest A")
+			void setForTestMethod() {
+				assertThat(System.getenv("set envvar A")).isEqualTo("newest A");
+			}
+
+			@AfterAll
+			@ReadsEnvironmentVariable
+			void resetAfterTestMethodExecution() {
+				assertThat(System.getenv("set envvar A")).isEqualTo("newer A");
+			}
+
+		}
+
 		@AfterAll
 		@ReadsEnvironmentVariable
 		void resetAfterTestContainerExecution() {
@@ -305,6 +336,15 @@ void shouldFailWhenSetSameEnvironmentVariableTwice() {
 			assertThat(results).hasSingleFailedTest().withExceptionInstanceOf(ExtensionConfigurationException.class);
 		}
 
+		@Test
+		@DisplayName("should fail when configuring test-level annotation with ApplyMode.CLASS")
+		void shouldFailWhenTestAnnotationIsAppliedOnClassLevel() {
+			ExecutionResults results = executeTestMethod(MethodLevelInitializationFailureTestCases.class,
+				"shouldFailWhenTestAnnotationIsAppliedOnClassLevel");
+
+			assertThat(results).hasSingleFailedTest().withExceptionInstanceOf(ExtensionConfigurationException.class);
+		}
+
 	}
 
 	@Nested
@@ -388,6 +428,11 @@ void shouldFailWhenClearSameEnvironmentVariableTwice() {
 		void shouldFailWhenSetSameEnvironmentVariableTwice() {
 		}
 
+		@Test
+		@SetEnvironmentVariable(key = "set envvar A", value = "new A", mode = ApplyMode.CLASS)
+		void shouldFailWhenTestAnnotationIsAppliedOnClassLevel() {
+		}
+
 	}
 
 	@Nested