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<byte[]> getTcpConnection() throws Exception {
-		WebSocketHandler webSocketHandler = connect();
-		webSocketHandler.afterConnectionEstablished(this.webSocketSession);
-		return (TcpConnection<byte[]>) webSocketHandler;
+		WebSocketHandler handler = connect();
+		handler.afterConnectionEstablished(this.webSocketSession);
+		if (handler instanceof WebSocketHandlerDecorator) {
+			handler = ((WebSocketHandlerDecorator) handler).getLastHandler();
+		}
+		return (TcpConnection<byte[]>) handler;
 	}
 
 	private void testInactivityTaskScheduling(Runnable runnable, long delay, long sleepTime)