diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/MessagingTextMap.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/MessagingTextMap.java index 88f79dfeb..49e67e164 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/MessagingTextMap.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/MessagingTextMap.java @@ -22,6 +22,7 @@ import java.util.Map; import org.springframework.cloud.sleuth.SpanTextMap; import org.springframework.messaging.Message; +import org.springframework.messaging.simp.SimpMessageHeaderAccessor; import org.springframework.messaging.support.MessageBuilder; import org.springframework.messaging.support.MessageHeaderAccessor; import org.springframework.messaging.support.NativeMessageHeaderAccessor; @@ -48,7 +49,9 @@ class MessagingTextMap implements SpanTextMap { Map map = new HashMap<>(); for (Map.Entry entry : this.delegate.build().getHeaders() .entrySet()) { - map.put(entry.getKey(), String.valueOf(entry.getValue())); + if (!NativeMessageHeaderAccessor.NATIVE_HEADERS.equals(entry.getKey())) { + map.put(entry.getKey(), String.valueOf(entry.getValue())); + } } return map.entrySet().iterator(); } @@ -63,8 +66,8 @@ class MessagingTextMap implements SpanTextMap { MessageHeaderAccessor accessor = MessageHeaderAccessor .getMutableAccessor(initialMessage); accessor.setHeader(key, value); - if (accessor instanceof NativeMessageHeaderAccessor) { - NativeMessageHeaderAccessor nativeAccessor = (NativeMessageHeaderAccessor) accessor; + if (accessor instanceof SimpMessageHeaderAccessor) { + SimpMessageHeaderAccessor nativeAccessor = (SimpMessageHeaderAccessor) accessor; nativeAccessor.setNativeHeader(key, value); } else if (accessor.getHeader(NativeMessageHeaderAccessor.NATIVE_HEADERS) != null) { diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjectorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjectorTests.java index 4ee13d5e4..a40d5501e 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjectorTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/HeaderBasedMessagingInjectorTests.java @@ -6,11 +6,12 @@ import java.util.Iterator; import java.util.Map; import org.junit.Test; + import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.SpanTextMap; import org.springframework.cloud.sleuth.TraceKeys; -import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.then; +import static org.assertj.core.api.BDDAssertions.then; /** * @author Marcin Grzejszczak @@ -19,7 +20,6 @@ public class HeaderBasedMessagingInjectorTests { HeaderBasedMessagingInjector injector = new HeaderBasedMessagingInjector(new TraceKeys()); - @SuppressWarnings("unchecked") @Test public void should_not_override_already_existing_headers() throws Exception { Span span = Span.builder() diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/MessagingTextMapTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/MessagingTextMapTests.java new file mode 100644 index 000000000..83f1c730a --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/MessagingTextMapTests.java @@ -0,0 +1,83 @@ +/* + * Copyright 2016-2017 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 + * + * http://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.cloud.sleuth.instrument.messaging; + +import java.util.HashSet; +import java.util.Set; + +import org.junit.Test; + +import org.springframework.messaging.Message; +import org.springframework.messaging.support.MessageBuilder; +import org.springframework.messaging.support.NativeMessageHeaderAccessor; +import org.springframework.util.LinkedMultiValueMap; +import org.springframework.util.MultiValueMap; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Dave Syer + * + */ +public class MessagingTextMapTests { + + @Test + public void vanilla() { + MessageBuilder builder = MessageBuilder.withPayload("foo"); + MessagingTextMap map = new MessagingTextMap(builder); + map.put("foo", "bar"); + Set keys = new HashSet<>(); + map.forEach(entry -> keys.add(entry.getKey())); + assertThat(keys).contains("foo"); + @SuppressWarnings("unchecked") + MultiValueMap natives = (MultiValueMap) builder.build().getHeaders().get(NativeMessageHeaderAccessor.NATIVE_HEADERS); + assertThat(natives).containsKey("foo"); + assertThat(keys).doesNotContain(NativeMessageHeaderAccessor.NATIVE_HEADERS); + } + + @Test + public void nativeHeadersAlreadyExist() { + MessageBuilder builder = MessageBuilder.withPayload("foo").setHeader( + NativeMessageHeaderAccessor.NATIVE_HEADERS, new LinkedMultiValueMap<>()); + MessagingTextMap map = new MessagingTextMap(builder); + map.put("foo", "bar"); + Set keys = new HashSet<>(); + map.forEach(entry -> keys.add(entry.getKey())); + assertThat(keys).contains("foo"); + @SuppressWarnings("unchecked") + MultiValueMap natives = (MultiValueMap) builder.build().getHeaders().get(NativeMessageHeaderAccessor.NATIVE_HEADERS); + assertThat(natives).containsKey("foo"); + assertThat(keys).doesNotContain(NativeMessageHeaderAccessor.NATIVE_HEADERS); + } + + @Test + public void nativeHeaders() { + Message message = MessageBuilder.withPayload("foo").build(); + MessageBuilder builder = MessageBuilder.fromMessage(message) + .setHeaders(NativeMessageHeaderAccessor.getMutableAccessor(message)); + MessagingTextMap map = new MessagingTextMap(builder); + map.put("foo", "bar"); + Set keys = new HashSet<>(); + map.forEach(entry -> keys.add(entry.getKey())); + assertThat(keys).contains("foo"); + @SuppressWarnings("unchecked") + MultiValueMap natives = (MultiValueMap) builder.build().getHeaders().get(NativeMessageHeaderAccessor.NATIVE_HEADERS); + assertThat(natives).containsKey("foo"); + assertThat(keys).doesNotContain(NativeMessageHeaderAccessor.NATIVE_HEADERS); + } + +} diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/websocket/TraceWebSocketAutoConfigurationTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/websocket/TraceWebSocketAutoConfigurationTests.java similarity index 96% rename from spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/websocket/TraceWebSocketAutoConfigurationTest.java rename to spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/websocket/TraceWebSocketAutoConfigurationTests.java index 442590476..d42d2587f 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/websocket/TraceWebSocketAutoConfigurationTest.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/websocket/TraceWebSocketAutoConfigurationTests.java @@ -39,8 +39,8 @@ import org.springframework.web.socket.config.annotation.StompEndpointRegistry; * @author Marcin Grzejszczak */ @RunWith(SpringRunner.class) -@SpringBootTest(classes = TraceWebSocketAutoConfigurationTest.Config.class) -public class TraceWebSocketAutoConfigurationTest { +@SpringBootTest(classes = TraceWebSocketAutoConfigurationTests.Config.class) +public class TraceWebSocketAutoConfigurationTests { @Autowired DelegatingWebSocketMessageBrokerConfiguration delegatingWebSocketMessageBrokerConfiguration;