CitrusWebSocketDispatcherServlet extends CitrusDispatcherServlet { private void configureWebSockerHandler(ApplicationContext context) { List<WebSocketEndpoint> webSocketEndpoints = webSocketServer.getWebSockets(); if (CollectionUtils.isEmpty(webSocketEndpoints)) { return; } if (context.containsBean(URL_HANDLER_MAPPING_BEAN_NAME)) { WebSocketUrlHandlerMapping urlHandlerMapping = context.getBean(URL_HANDLER_MAPPING_BEAN_NAME, WebSocketUrlHandlerMapping.class); HandshakeHandler handshakeHandler = new DefaultHandshakeHandler(); if (context.containsBean(HANDSHAKE_HANDLER_BEAN_NAME)) { handshakeHandler = context.getBean(HANDSHAKE_HANDLER_BEAN_NAME, HandshakeHandler.class); } Map<String, Object> wsHandlers = new HashMap<>(); for (WebSocketEndpoint webSocketEndpoint : webSocketEndpoints) { String wsPath = webSocketEndpoint.getEndpointConfiguration().getEndpointUri(); CitrusWebSocketHandler handler = new CitrusWebSocketHandler(); webSocketEndpoint.setWebSocketHandler(handler); WebSocketHttpRequestHandler wsRequestHandler = new WebSocketHttpRequestHandler(handler, handshakeHandler); SessionEnricherHandshakeInterceptor handshakeInterceptor = new SessionEnricherHandshakeInterceptor(webSocketEndpoint.getName(), wsPath); wsRequestHandler.getHandshakeInterceptors().add(handshakeInterceptor); wsHandlers.put(wsPath, wsRequestHandler); } urlHandlerMapping.setUrlMap(wsHandlers); urlHandlerMapping.postRegisterUrlHandlers(wsHandlers); } else { throw new CitrusRuntimeException(String.format("Invalid WebSocket configuration - missing bean. Expected to find bean with name '%s' in Spring application context", URL_HANDLER_MAPPING_BEAN_NAME)); } } CitrusWebSocketDispatcherServlet(WebSocketServer webSocketServer);   }
@Test public void testConfigureWebSockerHandler() throws Exception { WebSocketEndpoint wsEndpoint = Mockito.mock(WebSocketEndpoint.class); WebSocketEndpointConfiguration wsEndpointConfig = Mockito.mock(WebSocketEndpointConfiguration.class); String wsId = "wsId"; String endpointUri = "someEndpointUri"; List<WebSocketEndpoint> webSockets = new ArrayList<>(); webSockets.add(wsEndpoint); WebAppContext.Context servletContext = Mockito.mock(WebAppContext.Context.class); ContextHandler contextHandler = Mockito.mock(ContextHandler.class); DecoratedObjectFactory objectFactory = Mockito.mock(DecoratedObjectFactory.class); reset(httpServer, servletContext, contextHandler); when(httpServer.getInterceptors()).thenReturn(null); when(httpServer.getEndpointAdapter()).thenReturn(null); when(httpServer.getMessageConverter()).thenReturn(new HttpMessageConverter()); when(httpServer.getWebSockets()).thenReturn(webSockets); when(wsEndpoint.getEndpointConfiguration()).thenReturn(wsEndpointConfig); when(wsEndpoint.getName()).thenReturn(wsId); wsEndpoint.setWebSocketHandler(isA(CitrusWebSocketHandler.class)); when(wsEndpointConfig.getEndpointUri()).thenReturn(endpointUri); when(servletContext.getContextHandler()).thenReturn(contextHandler); when(servletContext.getAttribute(DecoratedObjectFactory.ATTR)).thenReturn(objectFactory); when(contextHandler.getServer()).thenReturn(new Server()); handshakeHandler.setServletContext(servletContext); servlet.initStrategies(applicationContext); Map<String, ?> urlMap = urlHandlerMapping.getUrlMap(); Assert.assertEquals(urlMap.size(), 1); Assert.assertTrue(urlMap.containsKey(endpointUri)); Assert.assertNotNull(urlMap.get(endpointUri)); }