Make entry-based extensions execute around classes

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

@@ -27,7 +27,9 @@
 import java.util.function.Function;
 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 +48,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, "Class", context));
+	}
 
 	@Override
 	public void beforeEach(ExtensionContext context) {
@@ -57,10 +66,10 @@ public void beforeEach(ExtensionContext context) {
 		 */
 		List<ExtensionContext> contexts = PioneerUtils.findAllContexts(context);
 		Collections.reverse(contexts);
-		contexts.forEach(this::clearAndSetEntries);
+		contexts.forEach(c -> clearAndSetEntries(c, "Test", context));
 	}
 
-	private void clearAndSetEntries(ExtensionContext context) {
+	private void clearAndSetEntries(ExtensionContext context, String mode, ExtensionContext original) {
 		context.getElement().ifPresent(element -> {
 			Set<K> entriesToClear;
 			Map<K, V> entriesToSet;
@@ -78,7 +87,7 @@ private void clearAndSetEntries(ExtensionContext context) {
 				return;
 
 			reportWarning(context);
-			storeOriginalEntries(context, entriesToClear, entriesToSet.keySet());
+			storeOriginalEntries(original, entriesToClear, entriesToSet.keySet(), mode);
 			clearEntries(entriesToClear);
 			setEntries(entriesToSet);
 		});
@@ -125,8 +134,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, String mode) {
+		getStore(context).put(getStoreKey(context, mode), new EntriesBackup(entriesToClear, entriesToSet));
 	}
 
 	private void clearEntries(Collection<K> entriesToClear) {
@@ -138,21 +147,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, "Test", context));
+	}
+
+	@Override
+	public void afterAll(ExtensionContext context) {
+		PioneerUtils.findAllContexts(context).forEach(c -> restoreOriginalEntries(c, "Class", context));
 	}
 
-	private void restoreOriginalEntries(ExtensionContext context) {
-		getStore(context).getOrDefault(getStoreKey(context), EntriesBackup.class, new EntriesBackup()).restoreBackup();
+	private void restoreOriginalEntries(ExtensionContext context, String mode, ExtensionContext original) {
+		getStore(context)
+				.getOrDefault(getStoreKey(original, 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, String mode) {
+		return context.getUniqueId() + mode;
 	}
 
 	private class EntriesBackup {

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

@@ -50,11 +50,8 @@ static void globalSetUp() {
 
 	@AfterAll
 	static void globalTearDown() {
-		assertThat(systemEnvironmentVariable("set envvar A")).isEqualTo("old A");
 		EnvironmentVariableUtils.clear("set envvar A");
-		assertThat(systemEnvironmentVariable("set envvar B")).isEqualTo("old B");
 		EnvironmentVariableUtils.clear("set envvar B");
-		assertThat(systemEnvironmentVariable("set envvar C")).isEqualTo("old C");
 		EnvironmentVariableUtils.clear("set envvar C");
 
 		assertThat(systemEnvironmentVariable("clear envvar D")).isNull();
@@ -247,7 +244,38 @@ class ResettingEnvironmentVariableTests {
 		@Nested
 		@SetEnvironmentVariable(key = "set envvar A", value = "newer A")
 		@TestInstance(TestInstance.Lifecycle.PER_CLASS)
-		class ResettingEnvironmentVariableNestedTests {
+		class ResettingEnvironmentVariableAfterEachNestedTests {
+
+			@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");
+			}
+
+		}
+
+		@Nested
+		@SetEnvironmentVariable(key = "set envvar A", value = "newer A")
+		@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")
@@ -258,15 +286,15 @@ void setForTestMethod() {
 			@AfterAll
 			@ReadsEnvironmentVariable
 			void resetAfterTestMethodExecution() {
-				assertThat(System.getenv("set envvar A")).isEqualTo("old A");
+				assertThat(System.getenv("set envvar A")).isEqualTo("newer A");
 			}
 
 		}
 
 		@AfterAll
 		@ReadsEnvironmentVariable
 		void resetAfterTestContainerExecution() {
-			assertThat(System.getenv("set envvar A")).isEqualTo("old A");
+			assertThat(System.getenv("set envvar A")).isEqualTo("new A");
 		}
 
 	}

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

@@ -11,6 +11,7 @@
 package org.junitpioneer.jupiter;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.junitpioneer.testkit.PioneerTestKit.executeTestMethod;
 import static org.junitpioneer.testkit.assertion.PioneerAssert.assertThat;
 
 import org.junit.jupiter.api.AfterAll;
@@ -25,7 +26,6 @@
 import org.junit.jupiter.api.TestMethodOrder;
 import org.junit.jupiter.api.extension.ExtensionConfigurationException;
 import org.junitpioneer.testkit.ExecutionResults;
-import org.junitpioneer.testkit.PioneerTestKit;
 
 @DisplayName("SystemProperty extension")
 class SystemPropertyExtensionTests {
@@ -43,11 +43,8 @@ static void globalSetUp() {
 
 	@AfterAll
 	static void globalTearDown() {
-		assertThat(System.getProperty("A")).isEqualTo("old A");
 		System.clearProperty("A");
-		assertThat(System.getProperty("B")).isEqualTo("old B");
 		System.clearProperty("B");
-		assertThat(System.getProperty("C")).isEqualTo("old C");
 		System.clearProperty("C");
 
 		assertThat(System.getProperty("clear prop D")).isNull();
@@ -236,7 +233,36 @@ class ResettingSystemPropertyTests {
 		@Nested
 		@SetSystemProperty(key = "A", value = "newer A")
 		@TestInstance(TestInstance.Lifecycle.PER_CLASS)
-		class ResettingSystemPropertyNestedTests {
+		class ResettingSystemPropertyAfterEachNestedTests {
+
+			@BeforeAll
+			void changeShouldBeVisible() {
+				assertThat(System.getProperty("A")).isEqualTo("newer A");
+			}
+
+			@Test
+			@SetSystemProperty(key = "A", value = "newest A")
+			void setForTestMethod() {
+				assertThat(System.getProperty("A")).isEqualTo("newest A");
+			}
+
+			@AfterAll
+			@ReadsSystemProperty
+			void resetAfterTestMethodExecution() {
+				assertThat(System.getProperty("A")).isEqualTo("newer A");
+			}
+
+		}
+
+		@Nested
+		@SetSystemProperty(key = "A", value = "newer A")
+		@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+		class ResettingSystemPropertyAfterAllNestedTests {
+
+			@BeforeAll
+			void changeShouldBeVisible() {
+				assertThat(System.getProperty("A")).isEqualTo("newer A");
+			}
 
 			@Test
 			@SetSystemProperty(key = "A", value = "newest A")
@@ -247,15 +273,15 @@ void setForTestMethod() {
 			@AfterAll
 			@ReadsSystemProperty
 			void resetAfterTestMethodExecution() {
-				assertThat(System.getProperty("A")).isEqualTo("old A");
+				assertThat(System.getProperty("A")).isEqualTo("newer A");
 			}
 
 		}
 
 		@AfterAll
 		@ReadsSystemProperty
 		void resetAfterTestContainerExecution() {
-			assertThat(System.getProperty("A")).isEqualTo("old A");
+			assertThat(System.getProperty("A")).isEqualTo("new A");
 		}
 
 	}
@@ -267,9 +293,8 @@ class ConfigurationFailureTests {
 		@Test
 		@DisplayName("should fail when clear and set same system property")
 		void shouldFailWhenClearAndSetSameSystemProperty() {
-			ExecutionResults results = PioneerTestKit
-					.executeTestMethod(MethodLevelInitializationFailureTestCases.class,
-						"shouldFailWhenClearAndSetSameSystemProperty");
+			ExecutionResults results = executeTestMethod(MethodLevelInitializationFailureTestCases.class,
+				"shouldFailWhenClearAndSetSameSystemProperty");
 
 			assertThat(results).hasSingleFailedTest().withExceptionInstanceOf(ExtensionConfigurationException.class);
 		}
@@ -280,19 +305,17 @@ void shouldFailWhenClearAndSetSameSystemProperty() {
 				+ "deduplicates identical annotations like the ones required for this test: "
 				+ "https://github.com/junit-team/junit5/issues/2131")
 		void shouldFailWhenClearSameSystemPropertyTwice() {
-			ExecutionResults results = PioneerTestKit
-					.executeTestMethod(MethodLevelInitializationFailureTestCases.class,
-						"shouldFailWhenClearSameSystemPropertyTwice");
+			ExecutionResults results = executeTestMethod(MethodLevelInitializationFailureTestCases.class,
+				"shouldFailWhenClearSameSystemPropertyTwice");
 
 			assertThat(results).hasSingleFailedTest().withExceptionInstanceOf(ExtensionConfigurationException.class);
 		}
 
 		@Test
 		@DisplayName("should fail when set same system property twice")
 		void shouldFailWhenSetSameSystemPropertyTwice() {
-			ExecutionResults results = PioneerTestKit
-					.executeTestMethod(MethodLevelInitializationFailureTestCases.class,
-						"shouldFailWhenSetSameSystemPropertyTwice");
+			ExecutionResults results = executeTestMethod(MethodLevelInitializationFailureTestCases.class,
+				"shouldFailWhenSetSameSystemPropertyTwice");
 			assertThat(results).hasSingleFailedTest().withExceptionInstanceOf(ExtensionConfigurationException.class);
 		}