Introduce TestExecutionListener for Micrometer's ObservationRegistry in the TestContext framework

framework-platform/framework-platform.gradle

@@ -8,7 +8,7 @@ javaPlatform {
 
 dependencies {
 	api(platform("com.fasterxml.jackson:jackson-bom:2.14.3"))
-	api(platform("io.micrometer:micrometer-bom:1.10.7"))
+	api(platform("io.micrometer:micrometer-bom:1.10.8"))
 	api(platform("io.netty:netty-bom:4.1.93.Final"))
 	api(platform("io.netty:netty5-bom:5.0.0.Alpha5"))
 	api(platform("io.projectreactor:reactor-bom:2022.0.7"))
@@ -46,7 +46,7 @@ dependencies {
 		api("commons-io:commons-io:2.11.0")
 		api("de.bechte.junit:junit-hierarchicalcontextrunner:4.12.2")
 		api("info.picocli:picocli:4.7.4")
-		api("io.micrometer:context-propagation:1.0.0")
+		api("io.micrometer:context-propagation:1.0.3")
 		api("io.mockk:mockk:1.13.4")
 		api("io.projectreactor.netty:reactor-netty5-http:2.0.0-M3")
 		api("io.projectreactor.tools:blockhound:1.0.8.RELEASE")

spring-test/spring-test.gradle

@@ -50,6 +50,8 @@ dependencies {
 	optional("io.projectreactor:reactor-test")
 	optional("org.jetbrains.kotlinx:kotlinx-coroutines-core")
 	optional("org.jetbrains.kotlinx:kotlinx-coroutines-reactor")
+	optional('io.micrometer:context-propagation')
+	optional('io.micrometer:micrometer-observation')
 	testImplementation(project(":spring-core-test"))
 	testImplementation(project(":spring-context-support"))
 	testImplementation(project(":spring-oxm"))
@@ -58,7 +60,6 @@ dependencies {
 	testImplementation(testFixtures(project(":spring-core")))
 	testImplementation(testFixtures(project(":spring-tx")))
 	testImplementation(testFixtures(project(":spring-web")))
-	testImplementation('io.micrometer:context-propagation')
 	testImplementation("jakarta.annotation:jakarta.annotation-api")
 	testImplementation("javax.cache:cache-api")
 	testImplementation("jakarta.ejb:jakarta.ejb-api")

spring-test/src/main/java/org/springframework/test/context/observation/MicrometerObservationThreadLocalTestExecutionListener.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2002-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.test.context.observation;
+
+import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
+
+import org.springframework.context.ApplicationContext;
+import org.springframework.core.Conventions;
+import org.springframework.test.context.TestContext;
+import org.springframework.test.context.TestExecutionListener;
+import org.springframework.test.context.support.AbstractTestExecutionListener;
+
+/**
+ * {@code ObservationThreadLocalTestExecutionListener} is an implementation of the {@link TestExecutionListener}
+ * SPI that updates the {@link ObservationThreadLocalAccessor} with the {@link ObservationRegistry}
+ * taken from the {@link ApplicationContext} present in the {@link TestContext}.
+ *
+ * <p>This implementation is not thread-safe.
+ *
+ * @author Marcin Grzejszczak
+ * @since 6.1
+ */
+public class MicrometerObservationThreadLocalTestExecutionListener extends AbstractTestExecutionListener {
+
+	/**
+	 * Attribute name for a {@link TestContext} attribute which contains the previously
+	 * set {@link ObservationRegistry} on the {@link ObservationThreadLocalAccessor}.
+	 * <p>After all tests from the current test class have completed, the previously stored {@link ObservationRegistry}
+	 * will be restored. If tests are ran concurrently this might cause issues
+	 * unless the {@link ObservationRegistry} is always the same (which should be the case most frequently).
+	 */
+	private static final String PREVIOUS_OBSERVATION_REGISTRY = Conventions.getQualifiedAttributeName(
+			MicrometerObservationThreadLocalTestExecutionListener.class, "previousObservationRegistry");
+
+	/**
+	 * Retrieves the current {@link ObservationRegistry} stored
+	 * on {@link ObservationThreadLocalAccessor} instance and stores it
+	 * in the {@link TestContext} attributes and overrides it with
+	 * one stored in {@link ApplicationContext} associated with
+	 * the {@link TestContext}.
+	 * @param testContext the test context for the test; never {@code null}
+	 */
+	@Override
+	public void beforeTestMethod(TestContext testContext) {
+		testContext.setAttribute(PREVIOUS_OBSERVATION_REGISTRY,
+				ObservationThreadLocalAccessor.getInstance().getObservationRegistry());
+		testContext.getApplicationContext()
+				.getBeanProvider(ObservationRegistry.class)
+				.ifAvailable(observationRegistry ->
+						ObservationThreadLocalAccessor.getInstance()
+								.setObservationRegistry(observationRegistry));
+	}
+
+	/**
+	 * Retrieves the previously stored {@link ObservationRegistry} and sets it back
+	 * on the {@link ObservationThreadLocalAccessor} instance.
+	 * @param testContext the test context for the test; never {@code null}
+	 */
+	@Override
+	public void afterTestMethod(TestContext testContext) {
+		ObservationRegistry previousObservationRegistry =
+				(ObservationRegistry) testContext.getAttribute(PREVIOUS_OBSERVATION_REGISTRY);
+		if (previousObservationRegistry != null) {
+			ObservationThreadLocalAccessor.getInstance()
+					.setObservationRegistry(previousObservationRegistry);
+		}
+	}
+
+
+	/**
+	 * Returns {@code 3500}.
+	 */
+	@Override
+	public final int getOrder() {
+		return 3500;
+	}
+}

spring-test/src/main/java/org/springframework/test/context/observation/package-info.java

@@ -0,0 +1,9 @@
+/**
+ * Observation support classes for the <em>Spring TestContext Framework</em>.
+ */
+@NonNullApi
+@NonNullFields
+package org.springframework.test.context.observation;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-test/src/main/resources/META-INF/spring.factories

@@ -8,7 +8,8 @@ org.springframework.test.context.TestExecutionListener = \
 	org.springframework.test.context.support.DirtiesContextTestExecutionListener,\
 	org.springframework.test.context.transaction.TransactionalTestExecutionListener,\
 	org.springframework.test.context.jdbc.SqlScriptsTestExecutionListener,\
-	org.springframework.test.context.event.EventPublishingTestExecutionListener
+	org.springframework.test.context.event.EventPublishingTestExecutionListener,\
+	org.springframework.test.context.observation.MicrometerObservationThreadLocalTestExecutionListener
 
 # Default ContextCustomizerFactory implementations for the Spring TestContext Framework
 #

spring-test/src/test/java/org/springframework/test/context/TestExecutionListenersTests.java

@@ -28,6 +28,7 @@
 import org.springframework.test.context.event.ApplicationEventsTestExecutionListener;
 import org.springframework.test.context.event.EventPublishingTestExecutionListener;
 import org.springframework.test.context.jdbc.SqlScriptsTestExecutionListener;
+import org.springframework.test.context.observation.MicrometerObservationThreadLocalTestExecutionListener;
 import org.springframework.test.context.support.AbstractTestExecutionListener;
 import org.springframework.test.context.support.DependencyInjectionTestExecutionListener;
 import org.springframework.test.context.support.DirtiesContextBeforeModesTestExecutionListener;
@@ -63,6 +64,7 @@ void defaultListeners() {
 				ApplicationEventsTestExecutionListener.class,//
 				DependencyInjectionTestExecutionListener.class,//
 				DirtiesContextTestExecutionListener.class,//
+				MicrometerObservationThreadLocalTestExecutionListener.class,//
 				TransactionalTestExecutionListener.class,//
 				SqlScriptsTestExecutionListener.class,//
 				EventPublishingTestExecutionListener.class
@@ -81,6 +83,7 @@ void defaultListenersMergedWithCustomListenerPrepended() {
 				ApplicationEventsTestExecutionListener.class,//
 				DependencyInjectionTestExecutionListener.class,//
 				DirtiesContextTestExecutionListener.class,//
+				MicrometerObservationThreadLocalTestExecutionListener.class,//
 				TransactionalTestExecutionListener.class,//
 				SqlScriptsTestExecutionListener.class,//
 				EventPublishingTestExecutionListener.class
@@ -98,6 +101,7 @@ void defaultListenersMergedWithCustomListenerAppended() {
 				ApplicationEventsTestExecutionListener.class,//
 				DependencyInjectionTestExecutionListener.class,//
 				DirtiesContextTestExecutionListener.class,//
+				MicrometerObservationThreadLocalTestExecutionListener.class,//
 				TransactionalTestExecutionListener.class,
 				SqlScriptsTestExecutionListener.class,//
 				EventPublishingTestExecutionListener.class,//
@@ -117,6 +121,7 @@ void defaultListenersMergedWithCustomListenerInserted() {
 				DependencyInjectionTestExecutionListener.class,//
 				BarTestExecutionListener.class,//
 				DirtiesContextTestExecutionListener.class,//
+				MicrometerObservationThreadLocalTestExecutionListener.class,//
 				TransactionalTestExecutionListener.class,//
 				SqlScriptsTestExecutionListener.class,//
 				EventPublishingTestExecutionListener.class

spring-test/src/test/java/org/springframework/test/context/observation/MicrometerObservationThreadLocalTestExecutionListenerTests.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2002-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.test.context.observation;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.context.support.StaticApplicationContext;
+import org.springframework.test.context.TestContext;
+
+import static org.assertj.core.api.BDDAssertions.then;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.BDDMockito.willAnswer;
+import static org.mockito.Mockito.mock;
+
+class MicrometerObservationThreadLocalTestExecutionListenerTests {
+
+	ObservationRegistry originalObservationRegistry = ObservationThreadLocalAccessor.getInstance().getObservationRegistry();
+
+	TestContext testContext = mock();
+
+	StaticApplicationContext applicationContext = new StaticApplicationContext();
+
+	Map<String, Object> attributes = new HashMap<>();
+
+	MicrometerObservationThreadLocalTestExecutionListener listener = new MicrometerObservationThreadLocalTestExecutionListener();
+
+	@BeforeEach
+	void setup() {
+		willAnswer(invocation -> attributes.put(invocation.getArgument(0), invocation.getArgument(1))).given(testContext).setAttribute(anyString(), any());
+		given(testContext.getAttribute(anyString())).willAnswer(invocation -> attributes.get(invocation.getArgument(0, String.class)));
+		given(testContext.getApplicationContext()).willReturn(applicationContext);
+	}
+
+	@Test
+	void observationRegistryShouldNotBeOverridden() throws Exception {
+		listener.beforeTestMethod(testContext);
+		thenObservationRegistryOnOTLAIsSameAsOriginal();
+		listener.afterTestMethod(testContext);
+		thenObservationRegistryOnOTLAIsSameAsOriginal();
+	}
+
+	@Test
+	void observationRegistryOverriddenByBeanFromTestContext() throws Exception {
+		ObservationRegistry newObservationRegistry = ObservationRegistry.create();
+		applicationContext.getDefaultListableBeanFactory().registerSingleton("observationRegistry", newObservationRegistry);
+
+		listener.beforeTestMethod(testContext);
+		ObservationRegistry otlaObservationRegistry = ObservationThreadLocalAccessor.getInstance().getObservationRegistry();
+		then(otlaObservationRegistry)
+				.as("During the test we want the original ObservationRegistry to be replaced with the one present in this application context")
+				.isNotSameAs(originalObservationRegistry)
+				.isSameAs(newObservationRegistry);
+
+		listener.afterTestMethod(testContext);
+		thenObservationRegistryOnOTLAIsSameAsOriginal();
+	}
+
+	private void thenObservationRegistryOnOTLAIsSameAsOriginal() {
+		then(ObservationThreadLocalAccessor.getInstance().getObservationRegistry()).isSameAs(originalObservationRegistry);
+	}
+
+}

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ReactiveTypeHandler.java

@@ -546,6 +546,7 @@ public ResolvableType getReturnType() {
 
 	private static class ContextSnapshotHelper {
 
+		@SuppressWarnings("deprecation")
 		public static Object writeReactorContext(Object returnValue) {
 			if (Mono.class.isAssignableFrom(returnValue.getClass())) {
 				ContextSnapshot snapshot = ContextSnapshot.captureAll();