Skip to content

Commit

Permalink
Introduce TestExecutionListener for Micrometer ObservationRegistry
Browse files Browse the repository at this point in the history
Prior to this commit, there was no way to specify the
ObservationRegistry that is registered in the given test's
ApplicationContext as the one that should be used by Micrometer's
ObservationThreadLocalAccessor for context propagation.

This commit introduces a TestExecutionListener for Micrometer's
ObservationRegistry in the Spring TestContext Framework. Specifically,
this listener obtains the ObservationRegistry registered in the test's
ApplicationContext, stores it in ObservationThreadLocalAccessor for the
duration of each test method execution, and restores the original
ObservationRegistry in ObservationThreadLocalAccessor after each test.

Co-authored-by: Sam Brannen <sam@sambrannen.com>
See spring-projectsgh-30658
  • Loading branch information
marcingrzejszczak authored and mdeinum committed Jun 29, 2023
1 parent 751d0c3 commit 085503f
Show file tree
Hide file tree
Showing 6 changed files with 194 additions and 2 deletions.
3 changes: 2 additions & 1 deletion spring-test/spring-test.gradle
Expand Up @@ -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"))
Expand All @@ -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")
Expand Down
@@ -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;
}
}
@@ -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;
3 changes: 2 additions & 1 deletion spring-test/src/main/resources/META-INF/spring.factories
Expand Up @@ -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
#
Expand Down
Expand Up @@ -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;
Expand Down Expand Up @@ -63,6 +64,7 @@ void defaultListeners() {
ApplicationEventsTestExecutionListener.class,//
DependencyInjectionTestExecutionListener.class,//
DirtiesContextTestExecutionListener.class,//
MicrometerObservationThreadLocalTestExecutionListener.class,//
TransactionalTestExecutionListener.class,//
SqlScriptsTestExecutionListener.class,//
EventPublishingTestExecutionListener.class
Expand All @@ -81,6 +83,7 @@ void defaultListenersMergedWithCustomListenerPrepended() {
ApplicationEventsTestExecutionListener.class,//
DependencyInjectionTestExecutionListener.class,//
DirtiesContextTestExecutionListener.class,//
MicrometerObservationThreadLocalTestExecutionListener.class,//
TransactionalTestExecutionListener.class,//
SqlScriptsTestExecutionListener.class,//
EventPublishingTestExecutionListener.class
Expand All @@ -98,6 +101,7 @@ void defaultListenersMergedWithCustomListenerAppended() {
ApplicationEventsTestExecutionListener.class,//
DependencyInjectionTestExecutionListener.class,//
DirtiesContextTestExecutionListener.class,//
MicrometerObservationThreadLocalTestExecutionListener.class,//
TransactionalTestExecutionListener.class,
SqlScriptsTestExecutionListener.class,//
EventPublishingTestExecutionListener.class,//
Expand All @@ -117,6 +121,7 @@ void defaultListenersMergedWithCustomListenerInserted() {
DependencyInjectionTestExecutionListener.class,//
BarTestExecutionListener.class,//
DirtiesContextTestExecutionListener.class,//
MicrometerObservationThreadLocalTestExecutionListener.class,//
TransactionalTestExecutionListener.class,//
SqlScriptsTestExecutionListener.class,//
EventPublishingTestExecutionListener.class
Expand Down
@@ -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);
}

}

0 comments on commit 085503f

Please sign in to comment.