Fix failing test after previous commit

See gh-23793
This commit is contained in:
Rossen Stoyanchev 2019-10-30 07:03:51 +00:00
parent 03ac8e6b42
commit 3895d21b7d
1 changed files with 7 additions and 3 deletions

View File

@ -46,6 +46,7 @@ import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHandler; import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession; import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.client.WebSocketClient; 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.assertThat;
import static org.assertj.core.api.Assertions.assertThatIllegalStateException; import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
@ -319,9 +320,12 @@ public class WebSocketStompClientTests {
@SuppressWarnings("unchecked") @SuppressWarnings("unchecked")
private TcpConnection<byte[]> getTcpConnection() throws Exception { private TcpConnection<byte[]> getTcpConnection() throws Exception {
WebSocketHandler webSocketHandler = connect(); WebSocketHandler handler = connect();
webSocketHandler.afterConnectionEstablished(this.webSocketSession); handler.afterConnectionEstablished(this.webSocketSession);
return (TcpConnection<byte[]>) webSocketHandler; if (handler instanceof WebSocketHandlerDecorator) {
handler = ((WebSocketHandlerDecorator) handler).getLastHandler();
}
return (TcpConnection<byte[]>) handler;
} }
private void testInactivityTaskScheduling(Runnable runnable, long delay, long sleepTime) private void testInactivityTaskScheduling(Runnable runnable, long delay, long sleepTime)