org.springframework.web.socket.messaging.SubProtocolWebSocketHandler Java Examples

@Test
public void customChannels() {
	loadBeanDefinitions("websocket-config-broker-customchannels.xml");

	List<Class<? extends MessageHandler>> subscriberTypes =
			Arrays.<Class<? extends MessageHandler>>asList(SimpAnnotationMethodMessageHandler.class,
					UserDestinationMessageHandler.class, SimpleBrokerMessageHandler.class);

	testChannel("clientInboundChannel", subscriberTypes, 3);
	testExecutor("clientInboundChannel", 100, 200, 600);

	subscriberTypes = Collections.singletonList(SubProtocolWebSocketHandler.class);

	testChannel("clientOutboundChannel", subscriberTypes, 3);
	testExecutor("clientOutboundChannel", 101, 201, 601);

	subscriberTypes = Arrays.<Class<? extends MessageHandler>>asList(SimpleBrokerMessageHandler.class,
			UserDestinationMessageHandler.class);

	testChannel("brokerChannel", subscriberTypes, 1);
	testExecutor("brokerChannel", 102, 202, 602);
}
 

@Bean
public WebSocketMessageBrokerStats webSocketMessageBrokerStats() {
	AbstractBrokerMessageHandler relayBean = stompBrokerRelayMessageHandler();
	StompBrokerRelayMessageHandler brokerRelay = (relayBean instanceof StompBrokerRelayMessageHandler ?
			(StompBrokerRelayMessageHandler) relayBean : null);

	// Ensure STOMP endpoints are registered
	stompWebSocketHandlerMapping();

	WebSocketMessageBrokerStats stats = new WebSocketMessageBrokerStats();
	stats.setSubProtocolWebSocketHandler((SubProtocolWebSocketHandler) subProtocolWebSocketHandler());
	stats.setStompBrokerRelay(brokerRelay);
	stats.setInboundChannelExecutor(clientInboundChannelExecutor());
	stats.setOutboundChannelExecutor(clientOutboundChannelExecutor());
	stats.setSockJsTaskScheduler(messageBrokerTaskScheduler());
	return stats;
}
 

@Test
public void clientInboundChannelSendMessage() throws Exception {
	ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
	TestChannel channel = config.getBean("clientInboundChannel", TestChannel.class);
	SubProtocolWebSocketHandler webSocketHandler = config.getBean(SubProtocolWebSocketHandler.class);

	List<ChannelInterceptor> interceptors = channel.getInterceptors();
	assertEquals(ImmutableMessageChannelInterceptor.class, interceptors.get(interceptors.size()-1).getClass());

	TestWebSocketSession session = new TestWebSocketSession("s1");
	session.setOpen(true);
	webSocketHandler.afterConnectionEstablished(session);

	TextMessage textMessage = StompTextMessageBuilder.create(StompCommand.SEND).headers("destination:/foo").build();
	webSocketHandler.handleMessage(session, textMessage);

	Message<?> message = channel.messages.get(0);
	StompHeaderAccessor accessor = StompHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
	assertNotNull(accessor);
	assertFalse(accessor.isMutable());
	assertEquals(SimpMessageType.MESSAGE, accessor.getMessageType());
	assertEquals("/foo", accessor.getDestination());
}
 

@Test
public void testInterceptors(final MockTracer tracer) {
  final DelegatingWebSocketMessageBrokerConfiguration configuration = new DelegatingWebSocketMessageBrokerConfiguration();

  final AbstractSubscribableChannel inboundChannel = configuration.clientInboundChannel();
  inboundChannel.setBeanName("clientInboundChannel");

  final AbstractSubscribableChannel outboundChannel = configuration.clientOutboundChannel();
  outboundChannel.setBeanName("clientOutboundChannel");

  outboundChannel.subscribe(new SubProtocolWebSocketHandler(inboundChannel, outboundChannel));
  inboundChannel.subscribe(new SubProtocolWebSocketHandler(inboundChannel, outboundChannel));

  configuration.clientInboundChannelExecutor().initialize();
  configuration.clientOutboundChannelExecutor().initialize();

  final Map<String,Object> headers = Collections.<String,Object>singletonMap("simpMessageType", SimpMessageType.MESSAGE);
  final GenericMessage<String> message = new GenericMessage<>("test", headers);
  outboundChannel.send(message);
  inboundChannel.send(message);

  await().atMost(15, TimeUnit.SECONDS).until(TestUtil.reportedSpansSize(tracer), equalTo(2));
  assertEquals(2, tracer.finishedSpans().size());
}
 

@Test
public void clientInboundChannelSendMessage() throws Exception {
	ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
	TestChannel channel = config.getBean("clientInboundChannel", TestChannel.class);
	SubProtocolWebSocketHandler webSocketHandler = config.getBean(SubProtocolWebSocketHandler.class);

	List<ChannelInterceptor> interceptors = channel.getInterceptors();
	assertEquals(ImmutableMessageChannelInterceptor.class, interceptors.get(interceptors.size()-1).getClass());

	TestWebSocketSession session = new TestWebSocketSession("s1");
	session.setOpen(true);
	webSocketHandler.afterConnectionEstablished(session);

	webSocketHandler.handleMessage(session,
			StompTextMessageBuilder.create(StompCommand.SEND).headers("destination:/foo").build());

	Message<?> message = channel.messages.get(0);
	StompHeaderAccessor accessor = StompHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
	assertNotNull(accessor);
	assertFalse(accessor.isMutable());
	assertEquals(SimpMessageType.MESSAGE, accessor.getMessageType());
	assertEquals("/foo", accessor.getDestination());
}
 

@Bean
public WebSocketMessageBrokerStats webSocketMessageBrokerStats() {
	AbstractBrokerMessageHandler relayBean = stompBrokerRelayMessageHandler();

	// Ensure STOMP endpoints are registered
	stompWebSocketHandlerMapping();

	WebSocketMessageBrokerStats stats = new WebSocketMessageBrokerStats();
	stats.setSubProtocolWebSocketHandler((SubProtocolWebSocketHandler) subProtocolWebSocketHandler());
	if (relayBean instanceof StompBrokerRelayMessageHandler) {
		stats.setStompBrokerRelay((StompBrokerRelayMessageHandler) relayBean);
	}
	stats.setInboundChannelExecutor(clientInboundChannelExecutor());
	stats.setOutboundChannelExecutor(clientOutboundChannelExecutor());
	stats.setSockJsTaskScheduler(messageBrokerTaskScheduler());
	return stats;
}
 

@Bean
public WebSocketMessageBrokerStats webSocketMessageBrokerStats() {
	AbstractBrokerMessageHandler relayBean = stompBrokerRelayMessageHandler();

	// Ensure STOMP endpoints are registered
	stompWebSocketHandlerMapping();

	WebSocketMessageBrokerStats stats = new WebSocketMessageBrokerStats();
	stats.setSubProtocolWebSocketHandler((SubProtocolWebSocketHandler) subProtocolWebSocketHandler());
	if (relayBean instanceof StompBrokerRelayMessageHandler) {
		stats.setStompBrokerRelay((StompBrokerRelayMessageHandler) relayBean);
	}
	stats.setInboundChannelExecutor(clientInboundChannelExecutor());
	stats.setOutboundChannelExecutor(clientOutboundChannelExecutor());
	stats.setSockJsTaskScheduler(messageBrokerTaskScheduler());
	return stats;
}
 

@Test
public void clientInboundChannelSendMessage() throws Exception {
	ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
	TestChannel channel = config.getBean("clientInboundChannel", TestChannel.class);
	SubProtocolWebSocketHandler webSocketHandler = config.getBean(SubProtocolWebSocketHandler.class);

	List<ChannelInterceptor> interceptors = channel.getInterceptors();
	assertEquals(ImmutableMessageChannelInterceptor.class, interceptors.get(interceptors.size()-1).getClass());

	TestWebSocketSession session = new TestWebSocketSession("s1");
	session.setOpen(true);
	webSocketHandler.afterConnectionEstablished(session);

	webSocketHandler.handleMessage(session,
			StompTextMessageBuilder.create(StompCommand.SEND).headers("destination:/foo").build());

	Message<?> message = channel.messages.get(0);
	StompHeaderAccessor accessor = StompHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
	assertNotNull(accessor);
	assertFalse(accessor.isMutable());
	assertEquals(SimpMessageType.MESSAGE, accessor.getMessageType());
	assertEquals("/foo", accessor.getDestination());
}
 

@Test
public void webSocketHandler() {
	ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
	SubProtocolWebSocketHandler subWsHandler = config.getBean(SubProtocolWebSocketHandler.class);

	assertEquals(1024 * 1024, subWsHandler.getSendBufferSizeLimit());
	assertEquals(25 * 1000, subWsHandler.getSendTimeLimit());
	assertEquals(30 * 1000, subWsHandler.getTimeToFirstMessage());

	Map<String, SubProtocolHandler> handlerMap = subWsHandler.getProtocolHandlerMap();
	StompSubProtocolHandler protocolHandler = (StompSubProtocolHandler) handlerMap.get("v12.stomp");
	assertEquals(128 * 1024, protocolHandler.getMessageSizeLimit());
}
 

@Test
public void getSubProtocols() throws Exception {
	SubscribableChannel channel = mock(SubscribableChannel.class);
	SubProtocolWebSocketHandler handler = new SubProtocolWebSocketHandler(channel, channel);
	StompSubProtocolHandler stompHandler = new StompSubProtocolHandler();
	handler.addProtocolHandler(stompHandler);

	TaskScheduler scheduler = mock(TaskScheduler.class);
	DefaultSockJsService service = new DefaultSockJsService(scheduler);
	WebSocketServerSockJsSession session = new WebSocketServerSockJsSession("1", service, handler, null);
	SockJsWebSocketHandler sockJsHandler = new SockJsWebSocketHandler(service, handler, session);

	assertEquals(stompHandler.getSupportedProtocols(), sockJsHandler.getSubProtocols());
}
 

@Override
public Message<?> beforeHandle(Message<?> message, MessageChannel channel,
    MessageHandler handler) {
  if ((handler instanceof WebSocketAnnotationMethodMessageHandler ||
      handler instanceof SubProtocolWebSocketHandler) &&
      SimpMessageType.MESSAGE.equals(message.getHeaders().get(SIMP_MESSAGE_TYPE))) {
    Span span = message.getHeaders().get(OPENTRACING_SPAN, Span.class);
    Scope scope = tracer.scopeManager().activate(span);
    message = MessageBuilder.fromMessage(message)
        .setHeader(OPENTRACING_SCOPE, scope)
        .build();
  }
  return message;
}
 

@Override
public void afterMessageHandled(Message<?> message, MessageChannel channel,
    MessageHandler handler, Exception arg3) {
  if ((handler instanceof WebSocketAnnotationMethodMessageHandler ||
      handler instanceof SubProtocolWebSocketHandler) &&
      SimpMessageType.MESSAGE.equals(message.getHeaders().get(SIMP_MESSAGE_TYPE))) {
    message.getHeaders().get(OPENTRACING_SCOPE, Scope.class).close();
    message.getHeaders().get(OPENTRACING_SPAN, Span.class).finish();
  }
}
 

@Override
public Message<?> beforeHandle(final Message<?> message, final MessageChannel channel, final MessageHandler handler) {
  if ((!(handler instanceof WebSocketAnnotationMethodMessageHandler) && !(handler instanceof SubProtocolWebSocketHandler)) || !SimpMessageType.MESSAGE.equals(message.getHeaders().get(SIMP_MESSAGE_TYPE)))
    return message;

  final Span span = message.getHeaders().get(OPENTRACING_SPAN, Span.class);
  final Scope scope = tracer.scopeManager().activate(span);
  return MessageBuilder.fromMessage(message).setHeader(OPENTRACING_SCOPE, scope).build();
}
 

@Override
public void afterMessageHandled(final Message<?> message, final MessageChannel channel, final MessageHandler handler, final Exception arg3) {
  if ((handler instanceof WebSocketAnnotationMethodMessageHandler || handler instanceof SubProtocolWebSocketHandler) && SimpMessageType.MESSAGE.equals(message.getHeaders().get(SIMP_MESSAGE_TYPE))) {
    message.getHeaders().get(OPENTRACING_SCOPE, Scope.class).close();
    message.getHeaders().get(OPENTRACING_SPAN, Span.class).finish();
    // MessageHeaders are immutable
    // message.getHeaders().remove(OPENTRACING_SCOPE);
  }
}
 

@Test
public void customChannels() {
	loadBeanDefinitions("websocket-config-broker-customchannels.xml");

	SimpAnnotationMethodMessageHandler annotationMethodMessageHandler =
			this.appContext.getBean(SimpAnnotationMethodMessageHandler.class);

	Validator validator = annotationMethodMessageHandler.getValidator();
	assertNotNull(validator);
	assertSame(this.appContext.getBean("myValidator"), validator);
	assertThat(validator, Matchers.instanceOf(TestValidator.class));

	List<Class<? extends MessageHandler>> subscriberTypes = Arrays.asList(SimpAnnotationMethodMessageHandler.class,
			UserDestinationMessageHandler.class, SimpleBrokerMessageHandler.class);

	testChannel("clientInboundChannel", subscriberTypes, 3);
	testExecutor("clientInboundChannel", 100, 200, 600);

	subscriberTypes = Collections.singletonList(SubProtocolWebSocketHandler.class);

	testChannel("clientOutboundChannel", subscriberTypes, 3);
	testExecutor("clientOutboundChannel", 101, 201, 601);

	subscriberTypes = Arrays.asList(SimpleBrokerMessageHandler.class, UserDestinationMessageHandler.class);

	testChannel("brokerChannel", subscriberTypes, 1);
	testExecutor("brokerChannel", 102, 202, 602);
}
 

@Test
public void clientOutboundChannel() {
	ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
	TestChannel channel = config.getBean("clientOutboundChannel", TestChannel.class);
	Set<MessageHandler> handlers = channel.getSubscribers();

	List<ChannelInterceptor> interceptors = channel.getInterceptors();
	assertEquals(ImmutableMessageChannelInterceptor.class, interceptors.get(interceptors.size()-1).getClass());

	assertEquals(1, handlers.size());
	assertTrue(handlers.contains(config.getBean(SubProtocolWebSocketHandler.class)));
}
 

@Before
public void setup() {
	SubscribableChannel inChannel = mock(SubscribableChannel.class);
	SubscribableChannel outChannel = mock(SubscribableChannel.class);
	this.webSocketHandler = new SubProtocolWebSocketHandler(inChannel, outChannel);

	WebSocketTransportRegistration transport = new WebSocketTransportRegistration();
	TaskScheduler scheduler = mock(TaskScheduler.class);
	this.endpointRegistry = new WebMvcStompEndpointRegistry(this.webSocketHandler, transport, scheduler);
}
 

private static SubProtocolWebSocketHandler unwrapSubProtocolWebSocketHandler(WebSocketHandler handler) {
	WebSocketHandler actual = WebSocketHandlerDecorator.unwrap(handler);
	if (!(actual instanceof SubProtocolWebSocketHandler)) {
		throw new IllegalArgumentException("No SubProtocolWebSocketHandler in " + handler);
	}
	return (SubProtocolWebSocketHandler) actual;
}
 

@Test
public void webSocketHandlerDecorator() throws Exception {
	ApplicationContext config = createConfig(WebSocketHandlerDecoratorConfig.class);
	WebSocketHandler handler = config.getBean(SubProtocolWebSocketHandler.class);
	assertNotNull(handler);

	SimpleUrlHandlerMapping mapping = (SimpleUrlHandlerMapping) config.getBean("stompWebSocketHandlerMapping");
	WebSocketHttpRequestHandler httpHandler = (WebSocketHttpRequestHandler) mapping.getHandlerMap().get("/test");
	handler = httpHandler.getWebSocketHandler();

	WebSocketSession session = new TestWebSocketSession("id");
	handler.afterConnectionEstablished(session);
	assertEquals(true, session.getAttributes().get("decorated"));
}
 

@Test
public void webSocketHandler() {
	ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
	SubProtocolWebSocketHandler subWsHandler = config.getBean(SubProtocolWebSocketHandler.class);

	assertEquals(1024 * 1024, subWsHandler.getSendBufferSizeLimit());
	assertEquals(25 * 1000, subWsHandler.getSendTimeLimit());
	assertEquals(30 * 1000, subWsHandler.getTimeToFirstMessage());

	Map<String, SubProtocolHandler> handlerMap = subWsHandler.getProtocolHandlerMap();
	StompSubProtocolHandler protocolHandler = (StompSubProtocolHandler) handlerMap.get("v12.stomp");
	assertEquals(128 * 1024, protocolHandler.getMessageSizeLimit());
}
 

@Test
public void clientOutboundChannel() {
	ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
	TestChannel channel = config.getBean("clientOutboundChannel", TestChannel.class);
	Set<MessageHandler> handlers = channel.getSubscribers();

	List<ChannelInterceptor> interceptors = channel.getInterceptors();
	assertEquals(ImmutableMessageChannelInterceptor.class, interceptors.get(interceptors.size()-1).getClass());

	assertEquals(1, handlers.size());
	assertTrue(handlers.contains(config.getBean(SubProtocolWebSocketHandler.class)));
}
 

@Test
public void getSubProtocols() throws Exception {
	SubscribableChannel channel = mock(SubscribableChannel.class);
	SubProtocolWebSocketHandler handler = new SubProtocolWebSocketHandler(channel, channel);
	StompSubProtocolHandler stompHandler = new StompSubProtocolHandler();
	handler.addProtocolHandler(stompHandler);

	TaskScheduler scheduler = mock(TaskScheduler.class);
	DefaultSockJsService service = new DefaultSockJsService(scheduler);
	WebSocketServerSockJsSession session = new WebSocketServerSockJsSession("1", service, handler, null);
	SockJsWebSocketHandler sockJsHandler = new SockJsWebSocketHandler(service, handler, session);

	assertEquals(stompHandler.getSupportedProtocols(), sockJsHandler.getSubProtocols());
}
 

@Test
public void getSubProtocols() throws Exception {
	SubscribableChannel channel = mock(SubscribableChannel.class);
	SubProtocolWebSocketHandler handler = new SubProtocolWebSocketHandler(channel, channel);
	StompSubProtocolHandler stompHandler = new StompSubProtocolHandler();
	handler.addProtocolHandler(stompHandler);

	TaskScheduler scheduler = mock(TaskScheduler.class);
	DefaultSockJsService service = new DefaultSockJsService(scheduler);
	WebSocketServerSockJsSession session = new WebSocketServerSockJsSession("1", service, handler, null);
	SockJsWebSocketHandler sockJsHandler = new SockJsWebSocketHandler(service, handler, session);

	assertEquals(stompHandler.getSupportedProtocols(), sockJsHandler.getSubProtocols());
}
 

@Test
public void webSocketHandler() {
	ApplicationContext config = createConfig(TestChannelConfig.class, TestConfigurer.class);
	SubProtocolWebSocketHandler subWsHandler = config.getBean(SubProtocolWebSocketHandler.class);

	assertEquals(1024 * 1024, subWsHandler.getSendBufferSizeLimit());
	assertEquals(25 * 1000, subWsHandler.getSendTimeLimit());

	Map<String, SubProtocolHandler> handlerMap = subWsHandler.getProtocolHandlerMap();
	StompSubProtocolHandler protocolHandler = (StompSubProtocolHandler) handlerMap.get("v12.stomp");
	assertEquals(128 * 1024, protocolHandler.getMessageSizeLimit());
}
 

@Test
public void webSocketHandlerDecorator() throws Exception {
	ApplicationContext config = createConfig(WebSocketHandlerDecoratorConfig.class);
	WebSocketHandler handler = config.getBean(SubProtocolWebSocketHandler.class);
	assertNotNull(handler);

	SimpleUrlHandlerMapping mapping = (SimpleUrlHandlerMapping) config.getBean("stompWebSocketHandlerMapping");
	WebSocketHttpRequestHandler httpHandler = (WebSocketHttpRequestHandler) mapping.getHandlerMap().get("/test");
	handler = httpHandler.getWebSocketHandler();

	WebSocketSession session = new TestWebSocketSession("id");
	handler.afterConnectionEstablished(session);
	assertEquals(true, session.getAttributes().get("decorated"));
}
 

private static SubProtocolWebSocketHandler unwrapSubProtocolWebSocketHandler(WebSocketHandler handler) {
	WebSocketHandler actual = WebSocketHandlerDecorator.unwrap(handler);
	if (!(actual instanceof SubProtocolWebSocketHandler)) {
		throw new IllegalArgumentException("No SubProtocolWebSocketHandler in " + handler);
	}
	return (SubProtocolWebSocketHandler) actual;
}
 

@Test
public void customChannels() {
	loadBeanDefinitions("websocket-config-broker-customchannels.xml");

	SimpAnnotationMethodMessageHandler annotationMethodMessageHandler =
			this.appContext.getBean(SimpAnnotationMethodMessageHandler.class);

	Validator validator = annotationMethodMessageHandler.getValidator();
	assertNotNull(validator);
	assertSame(this.appContext.getBean("myValidator"), validator);
	assertThat(validator, Matchers.instanceOf(TestValidator.class));

	List<Class<? extends MessageHandler>> subscriberTypes = Arrays.asList(SimpAnnotationMethodMessageHandler.class,
			UserDestinationMessageHandler.class, SimpleBrokerMessageHandler.class);

	testChannel("clientInboundChannel", subscriberTypes, 3);
	testExecutor("clientInboundChannel", 100, 200, 600);

	subscriberTypes = Collections.singletonList(SubProtocolWebSocketHandler.class);

	testChannel("clientOutboundChannel", subscriberTypes, 3);
	testExecutor("clientOutboundChannel", 101, 201, 601);

	subscriberTypes = Arrays.asList(SimpleBrokerMessageHandler.class, UserDestinationMessageHandler.class);

	testChannel("brokerChannel", subscriberTypes, 1);
	testExecutor("brokerChannel", 102, 202, 602);
}
 

@Before
public void setup() {
	SubscribableChannel inChannel = mock(SubscribableChannel.class);
	SubscribableChannel outChannel = mock(SubscribableChannel.class);
	this.webSocketHandler = new SubProtocolWebSocketHandler(inChannel, outChannel);

	WebSocketTransportRegistration transport = new WebSocketTransportRegistration();
	TaskScheduler scheduler = mock(TaskScheduler.class);
	this.endpointRegistry = new WebMvcStompEndpointRegistry(this.webSocketHandler, transport, scheduler);
}
 

@Before
public void setup() {
	SubscribableChannel inChannel = mock(SubscribableChannel.class);
	SubscribableChannel outChannel = mock(SubscribableChannel.class);
	this.webSocketHandler = new SubProtocolWebSocketHandler(inChannel, outChannel);

	WebSocketTransportRegistration transport = new WebSocketTransportRegistration();
	TaskScheduler scheduler = mock(TaskScheduler.class);
	this.endpointRegistry = new WebMvcStompEndpointRegistry(this.webSocketHandler, transport, null, scheduler);
}
 

@Test
public void webSocketHandlerDecorator() throws Exception {
	ApplicationContext config = createConfig(WebSocketHandlerDecoratorConfig.class);
	WebSocketHandler handler = config.getBean(SubProtocolWebSocketHandler.class);
	assertNotNull(handler);

	SimpleUrlHandlerMapping mapping = (SimpleUrlHandlerMapping) config.getBean("stompWebSocketHandlerMapping");
	WebSocketHttpRequestHandler httpHandler = (WebSocketHttpRequestHandler) mapping.getHandlerMap().get("/test");
	handler = httpHandler.getWebSocketHandler();

	WebSocketSession session = new TestWebSocketSession("id");
	handler.afterConnectionEstablished(session);
	assertEquals(true, session.getAttributes().get("decorated"));
}