From cc84533d85fce3cb33a9ef4cba016434eebbb4b8 Mon Sep 17 00:00:00 2001 From: Rossen Stoyanchev Date: Wed, 30 Oct 2019 06:53:25 +0000 Subject: [PATCH] Logging decorator for WebSocketStompClient handler Closes gh-23793 --- .../web/socket/messaging/WebSocketStompClient.java | 7 +++++-- .../socket/messaging/WebSocketStompClientTests.java | 10 +++++++--- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/WebSocketStompClient.java b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/WebSocketStompClient.java index dfdcc1830b..59507459fd 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/WebSocketStompClient.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/WebSocketStompClient.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2018 the original author or authors. + * Copyright 2002-2019 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. @@ -57,6 +57,7 @@ import org.springframework.web.socket.WebSocketHttpHeaders; import org.springframework.web.socket.WebSocketMessage; import org.springframework.web.socket.WebSocketSession; import org.springframework.web.socket.client.WebSocketClient; +import org.springframework.web.socket.handler.LoggingWebSocketHandlerDecorator; import org.springframework.web.socket.sockjs.transport.SockJsSession; import org.springframework.web.util.UriComponentsBuilder; @@ -265,7 +266,9 @@ public class WebSocketStompClient extends StompClientSupport implements SmartLif Assert.notNull(url, "'url' must not be null"); ConnectionHandlingStompSession session = createSession(connectHeaders, sessionHandler); WebSocketTcpConnectionHandlerAdapter adapter = new WebSocketTcpConnectionHandlerAdapter(session); - getWebSocketClient().doHandshake(adapter, handshakeHeaders, url).addCallback(adapter); + getWebSocketClient() + .doHandshake(new LoggingWebSocketHandlerDecorator(adapter), handshakeHeaders, url) + .addCallback(adapter); return session.getSessionFuture(); } diff --git a/spring-websocket/src/test/java/org/springframework/web/socket/messaging/WebSocketStompClientTests.java b/spring-websocket/src/test/java/org/springframework/web/socket/messaging/WebSocketStompClientTests.java index cf98cf0f0b..026f535d50 100644 --- a/spring-websocket/src/test/java/org/springframework/web/socket/messaging/WebSocketStompClientTests.java +++ b/spring-websocket/src/test/java/org/springframework/web/socket/messaging/WebSocketStompClientTests.java @@ -45,6 +45,7 @@ import org.springframework.web.socket.TextMessage; import org.springframework.web.socket.WebSocketHandler; import org.springframework.web.socket.WebSocketSession; import org.springframework.web.socket.client.WebSocketClient; +import org.springframework.web.socket.handler.WebSocketHandlerDecorator; import static org.junit.Assert.*; import static org.mockito.Mockito.*; @@ -317,9 +318,12 @@ public class WebSocketStompClientTests { @SuppressWarnings("unchecked") private TcpConnection getTcpConnection() throws Exception { - WebSocketHandler webSocketHandler = connect(); - webSocketHandler.afterConnectionEstablished(this.webSocketSession); - return (TcpConnection) webSocketHandler; + WebSocketHandler handler = connect(); + handler.afterConnectionEstablished(this.webSocketSession); + if (handler instanceof WebSocketHandlerDecorator) { + handler = ((WebSocketHandlerDecorator) handler).getLastHandler(); + } + return (TcpConnection) handler; } private void testInactivityTaskScheduling(Runnable runnable, long delay, long sleepTime)