Skip to content

Commit c5de98a

Browse files
rstoyanchevpull[bot]
authored andcommitted
Fix failing test after previous commit
See spring-projectsgh-23793
1 parent 473ce95 commit c5de98a

File tree

1 file changed

+7
-3
lines changed

1 file changed

+7
-3
lines changed

spring-websocket/src/test/java/org/springframework/web/socket/messaging/WebSocketStompClientTests.java

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import org.springframework.web.socket.WebSocketHandler;
4747
import org.springframework.web.socket.WebSocketSession;
4848
import org.springframework.web.socket.client.WebSocketClient;
49+
import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
4950

5051
import static org.assertj.core.api.Assertions.assertThat;
5152
import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
@@ -319,9 +320,12 @@ private WebSocketHandler connect() {
319320

320321
@SuppressWarnings("unchecked")
321322
private TcpConnection<byte[]> getTcpConnection() throws Exception {
322-
WebSocketHandler webSocketHandler = connect();
323-
webSocketHandler.afterConnectionEstablished(this.webSocketSession);
324-
return (TcpConnection<byte[]>) webSocketHandler;
323+
WebSocketHandler handler = connect();
324+
handler.afterConnectionEstablished(this.webSocketSession);
325+
if (handler instanceof WebSocketHandlerDecorator) {
326+
handler = ((WebSocketHandlerDecorator) handler).getLastHandler();
327+
}
328+
return (TcpConnection<byte[]>) handler;
325329
}
326330

327331
private void testInactivityTaskScheduling(Runnable runnable, long delay, long sleepTime)

0 commit comments

Comments
 (0)