From 3895d21b7ddc5ce378d31ca4befdc83450e16521 Mon Sep 17 00:00:00 2001 From: Rossen Stoyanchev Date: Wed, 30 Oct 2019 07:03:51 +0000 Subject: [PATCH] Fix failing test after previous commit See gh-23793 --- .../socket/messaging/WebSocketStompClientTests.java | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) 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 f65ef7ed0c0..4d849bde949 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 @@ -46,6 +46,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.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatIllegalStateException; @@ -319,9 +320,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)