Added more tests for websockets

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/websocket/TraceWebSocketAutoConfiguration.java

@@ -35,6 +35,9 @@ public class TraceWebSocketAutoConfiguration
 	@Autowired
 	private TraceKeys traceKeys;
 
+	@Autowired
+	private Random random;
+
 	@Override
 	public void registerStompEndpoints(StompEndpointRegistry registry) {
 		// The user must register their own endpoints
@@ -43,12 +46,12 @@ public class TraceWebSocketAutoConfiguration
 	@Override
 	public void configureClientOutboundChannel(ChannelRegistration registration) {
 		registration.setInterceptors(
-				new TraceChannelInterceptor(this.tracer, this.traceKeys, new Random()));
+				new TraceChannelInterceptor(this.tracer, this.traceKeys, this.random));
 	}
 
 	@Override
 	public void configureClientInboundChannel(ChannelRegistration registration) {
 		registration.setInterceptors(
-				new TraceChannelInterceptor(this.tracer, this.traceKeys, new Random()));
+				new TraceChannelInterceptor(this.tracer, this.traceKeys, this.random));
 	}
 }

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/websocket/TraceWebSocketAutoConfigurationTest.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2013-2016 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.websocket;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.test.SpringApplicationConfiguration;
+import org.springframework.cloud.sleuth.instrument.messaging.TraceChannelInterceptor;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.messaging.simp.config.MessageBrokerRegistry;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.web.socket.config.annotation.AbstractWebSocketMessageBrokerConfigurer;
+import org.springframework.web.socket.config.annotation.DelegatingWebSocketMessageBrokerConfiguration;
+import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
+import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
+
+import static java.util.stream.Collectors.toList;
+import static org.assertj.core.api.BDDAssertions.then;
+
+/**
+ * @author Marcin Grzejszczak
+ */
+@RunWith(SpringJUnit4ClassRunner.class)
+@SpringApplicationConfiguration(classes = TraceWebSocketAutoConfigurationTest.Config.class)
+public class TraceWebSocketAutoConfigurationTest {
+
+	@Autowired DelegatingWebSocketMessageBrokerConfiguration delegatingWebSocketMessageBrokerConfiguration;
+
+	@Test
+	public void should_register_interceptors_for_inbound_and_outbound_channels() {
+		then(delegatingWebSocketMessageBrokerConfiguration
+					.clientInboundChannel()
+					.getInterceptors()
+					.stream()
+					.map(Object::getClass)
+					.collect(toList()))
+				.contains(TraceChannelInterceptor.class);
+		then(delegatingWebSocketMessageBrokerConfiguration
+					.clientOutboundChannel()
+					.getInterceptors()
+					.stream()
+					.map(Object::getClass)
+					.collect(toList()))
+				.contains(TraceChannelInterceptor.class);
+	}
+
+	@EnableAutoConfiguration
+	@Configuration
+	@EnableWebSocketMessageBroker
+	public static class Config extends AbstractWebSocketMessageBrokerConfigurer {
+
+		@Override
+		public void configureMessageBroker(MessageBrokerRegistry config) {
+			config.enableSimpleBroker("/topic");
+			config.setApplicationDestinationPrefixes("/app");
+		}
+
+		@Override
+		public void registerStompEndpoints(StompEndpointRegistry registry) {
+			registry.addEndpoint("/hello").withSockJS();
+		}
+	}
+}