Make DefaultLocaleExtension work for class-level annotations

Author: marcphilipp

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/extension/DefaultLocaleExtension.java

@@ -10,10 +10,11 @@
 
 package org.junit.jupiter.engine.extension;
 
-import java.lang.reflect.AnnotatedElement;
 import java.util.Locale;
+import java.util.Optional;
 
 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;
@@ -25,30 +26,34 @@
 import org.junit.platform.commons.support.AnnotationSupport;
 import org.junit.platform.commons.support.ReflectionSupport;
 
-class DefaultLocaleExtension implements BeforeEachCallback, AfterEachCallback {
+class DefaultLocaleExtension implements BeforeAllCallback, BeforeEachCallback, AfterEachCallback {
 
 	private static final Namespace NAMESPACE = Namespace.create(DefaultLocaleExtension.class);
 
-	private static final String KEY = "DefaultLocale";
+	private static final String CUSTOM_KEY = "CustomLocale";
+	private static final String DEFAULT_KEY = "DefaultLocale";
+
+	@Override
+	public void beforeAll(ExtensionContext context) throws Exception {
+		createLocaleFromAnnotation(context) //
+				.ifPresent(locale -> store(context, CUSTOM_KEY, locale));
+	}
 
 	@Override
 	public void beforeEach(ExtensionContext context) {
-		AnnotatedElement element = context.getElement().orElse(null);
-		AnnotationSupport.findAnnotation(element, DefaultLocale.class).ifPresent(
-			annotation -> setDefaultLocale(context, annotation));
+		createLocaleFromAnnotation(context) //
+				.or(() -> load(context, CUSTOM_KEY)) //
+				.ifPresent(locale -> setDefaultLocale(context, locale));
 	}
 
-	private void setDefaultLocale(ExtensionContext context, DefaultLocale annotation) {
-		Locale configuredLocale = createLocale(annotation);
-		// defer storing the current default locale until the new locale could be created from the configuration
-		// (this prevents cases where misconfigured extensions store default locale now and restore it later,
-		// which leads to race conditions in our tests)
-		storeDefaultLocale(context);
-		Locale.setDefault(configuredLocale);
+	private void setDefaultLocale(ExtensionContext context, Locale customLocale) {
+		store(context, DEFAULT_KEY, Locale.getDefault());
+		Locale.setDefault(customLocale);
 	}
 
-	private void storeDefaultLocale(ExtensionContext context) {
-		context.getStore(NAMESPACE).put(KEY, Locale.getDefault());
+	private static Optional<Locale> createLocaleFromAnnotation(ExtensionContext context) {
+		return AnnotationSupport.findAnnotation(context.getElement(), DefaultLocale.class) //
+				.map(DefaultLocaleExtension::createLocale);
 	}
 
 	private static Locale createLocale(DefaultLocale annotation) {
@@ -108,23 +113,33 @@ private static Locale getFromProvider(DefaultLocale annotation) {
 			throw new ExtensionConfigurationException(
 				"LocaleProvider instance could not be constructed because of an exception", exception);
 		}
+		return invoke(provider);
+	}
+
+	@SuppressWarnings("ConstantValue")
+	private static Locale invoke(LocaleProvider provider) {
 		var locale = provider.get();
-		if (locale == null)
-			throw new NullPointerException("LocaleProvider instance returned with null");
+		if (locale == null) {
+			throw new ExtensionConfigurationException("LocaleProvider instance returned with null");
+		}
 		return locale;
 	}
 
 	@Override
 	public void afterEach(ExtensionContext context) {
-		AnnotatedElement element = context.getElement().orElse(null);
-		AnnotationSupport.findAnnotation(element, DefaultLocale.class).ifPresent(__ -> resetDefaultLocale(context));
+		load(context, DEFAULT_KEY).ifPresent(Locale::setDefault);
+	}
+
+	private static void store(ExtensionContext context, String key, Locale value) {
+		getStore(context).put(key, value);
+	}
+
+	private static Optional<Locale> load(ExtensionContext context, String key) {
+		return Optional.ofNullable(getStore(context).get(key, Locale.class));
 	}
 
-	private void resetDefaultLocale(ExtensionContext context) {
-		Locale defaultLocale = context.getStore(NAMESPACE).get(KEY, Locale.class);
-		// default locale is null if the extension was misconfigured and execution failed in "before"
-		if (defaultLocale != null)
-			Locale.setDefault(defaultLocale);
+	private static ExtensionContext.Store getStore(ExtensionContext context) {
+		return context.getStore(NAMESPACE);
 	}
 
 }