org.springframework.util.concurrent.SettableListenableFuture

/**
 * Unit tests for {@link DefaultTransportRequest}.
 *
 * @author Rossen Stoyanchev
 */
public clreplaced DefaultTransportRequestTests {

    private static final Jackson2SockJsMessageCodec CODEC = new Jackson2SockJsMessageCodec();

    private SettableListenableFuture<WebSocketSession> connectFuture;

    private ListenableFutureCallback<WebSocketSession> connectCallback;

    private TestTransport webSocketTransport;

    private TestTransport xhrTransport;

    @SuppressWarnings("unchecked")
    @Before
    public void setup() throws Exception {
        this.connectCallback = mock(ListenableFutureCallback.clreplaced);
        this.connectFuture = new SettableListenableFuture<>();
        this.connectFuture.addCallback(this.connectCallback);
        this.webSocketTransport = new TestTransport("WebSocketTestTransport");
        this.xhrTransport = new TestTransport("XhrTestTransport");
    }

    @Test
    public void connect() throws Exception {
        DefaultTransportRequest request = createTransportRequest(this.webSocketTransport, TransportType.WEBSOCKET);
        request.connect(null, this.connectFuture);
        WebSocketSession session = mock(WebSocketSession.clreplaced);
        this.webSocketTransport.getConnectCallback().onSuccess(session);
        replacedertSame(session, this.connectFuture.get());
    }

    @Test
    public void fallbackAfterTransportError() throws Exception {
        DefaultTransportRequest request1 = createTransportRequest(this.webSocketTransport, TransportType.WEBSOCKET);
        DefaultTransportRequest request2 = createTransportRequest(this.xhrTransport, TransportType.XHR_STREAMING);
        request1.setFallbackRequest(request2);
        request1.connect(null, this.connectFuture);
        // Transport error => fallback
        this.webSocketTransport.getConnectCallback().onFailure(new IOException("Fake exception 1"));
        replacedertFalse(this.connectFuture.isDone());
        replacedertTrue(this.xhrTransport.invoked());
        // Transport error => no more fallback
        this.xhrTransport.getConnectCallback().onFailure(new IOException("Fake exception 2"));
        replacedertTrue(this.connectFuture.isDone());
        replacedertThatExceptionOfType(ExecutionException.clreplaced).isThrownBy(this.connectFuture::get).withMessageContaining("Fake exception 2");
    }

    @Test
    public void fallbackAfterTimeout() throws Exception {
        TaskScheduler scheduler = mock(TaskScheduler.clreplaced);
        Runnable sessionCleanupTask = mock(Runnable.clreplaced);
        DefaultTransportRequest request1 = createTransportRequest(this.webSocketTransport, TransportType.WEBSOCKET);
        DefaultTransportRequest request2 = createTransportRequest(this.xhrTransport, TransportType.XHR_STREAMING);
        request1.setFallbackRequest(request2);
        request1.setTimeoutScheduler(scheduler);
        request1.addTimeoutTask(sessionCleanupTask);
        request1.connect(null, this.connectFuture);
        replacedertTrue(this.webSocketTransport.invoked());
        replacedertFalse(this.xhrTransport.invoked());
        // Get and invoke the scheduled timeout task
        ArgumentCaptor<Runnable> taskCaptor = ArgumentCaptor.forClreplaced(Runnable.clreplaced);
        verify(scheduler).schedule(taskCaptor.capture(), any(Date.clreplaced));
        verifyNoMoreInteractions(scheduler);
        taskCaptor.getValue().run();
        replacedertTrue(this.xhrTransport.invoked());
        verify(sessionCleanupTask).run();
    }

    protected DefaultTransportRequest createTransportRequest(Transport transport, TransportType type) throws Exception {
        SockJsUrlInfo urlInfo = new SockJsUrlInfo(new URI("https://example.com"));
        return new DefaultTransportRequest(urlInfo, new HttpHeaders(), new HttpHeaders(), transport, type, CODEC);
    }
}

@Test
public void sendWithExecutionException() {
    this.session.afterConnected(this.connection);
    replacedertTrue(this.session.isConnected());
    IllegalStateException exception = new IllegalStateException("simulated exception");
    SettableListenableFuture<Void> future = new SettableListenableFuture<>();
    future.setException(exception);
    given(this.connection.send(any())).willReturn(future);
    replacedertThatExceptionOfType(MessageDeliveryException.clreplaced).isThrownBy(() -> this.session.send("/topic/foo", "sample payload".getBytes(StandardCharsets.UTF_8))).withCause(exception);
}

private Release handleMessageAndWait(Message<SkipperEvents> message, String machineId, SkipperStates... statesToWait) {
    // machine gets acquired fully started
    StateMachine<SkipperStates, SkipperEvents> stateMachine = stateMachineService.acquireStateMachine(machineId);
    // setup future handling blocking requirement returning release
    SettableListenableFuture<Release> future = new SettableListenableFuture<>();
    StateMachineListener<SkipperStates, SkipperEvents> listener = new StateMachineListenerAdapter<SkipperStates, SkipperEvents>() {

        @Override
        public void stateContext(StateContext<SkipperStates, SkipperEvents> stateContext) {
            if (stateContext.getStage() == Stage.STATE_ENTRY) {
                if (stateContext.getTarget().getId() == SkipperStates.ERROR) {
                    Exception exception = stateContext.getExtendedState().get(SkipperVariables.ERROR, Exception.clreplaced);
                    if (exception != null) {
                        // we went through error state, throw if there is an error
                        log.info("setting future exception", exception);
                        future.setException(exception);
                    }
                } else if (Arrays.asList(statesToWait).contains(stateContext.getTarget().getId()) && !isInitialTransition(stateContext.getTransition())) {
                    Release release = (Release) stateContext.getExtendedState().getVariables().get(SkipperVariables.RELEASE);
                    // at this point we replacedume machine logic did set release
                    log.info("setting future value {}", release);
                    future.set(release);
                }
            }
        }
    };
    // add listener which gets removed eventually
    stateMachine.addStateListener(listener);
    future.addCallback(result -> {
        stateMachine.removeStateListener(listener);
    }, throwable -> {
        stateMachine.removeStateListener(listener);
    });
    // if machine doesn't accept an event, we're on state
    // where a particular message cannot be handled, thus
    // return exception. this simply happens when we are
    // i.e. upgrading and delete request comes in.
    if (stateMachine.sendEvent(message)) {
        try {
            return future.get();
        } catch (ExecutionException e) {
            if (e.getCause() instanceof SkipperException) {
                // throw as SkipperException
                throw (SkipperException) e.getCause();
            }
            throw new SkipperException("Error waiting to get Release from a statemachine", e);
        } catch (Exception e) {
            throw new SkipperException("Error waiting to get Release from a statemachine", e);
        }
    } else {
        throw new SkipperException("Statemachine is not in state ready to do " + message.getPayload());
    }
}

private ListenableFuture<DoreplacedentOcrResultSet> extractOcrResultFuture(OperationFuture<AsyncBatchAnnotateFilesResponse, OperationMetadata> grpcFuture) {
    SettableListenableFuture<DoreplacedentOcrResultSet> result = new SettableListenableFuture<>();
    ApiFutures.addCallback(grpcFuture, new ApiFutureCallback<AsyncBatchAnnotateFilesResponse>() {

        @Override
        public void onFailure(Throwable throwable) {
            result.setException(throwable);
        }

        @Override
        public void onSuccess(AsyncBatchAnnotateFilesResponse asyncBatchAnnotateFilesResponse) {
            String outputLocationUri = asyncBatchAnnotateFilesResponse.getResponsesList().get(0).getOutputConfig().getGcsDestination().getUri();
            GoogleStorageLocation outputFolderLocation = new GoogleStorageLocation(outputLocationUri);
            result.set(readOcrOutputFileSet(outputFolderLocation));
        }
    }, this.executor);
    return result;
}

/**
 * Perform Pub/Sub operations (ack/nack/modifyAckDeadline) in per-subscription batches.
 * <p>The returned {@link ListenableFuture} will complete when either all batches completes successfully or when at
 * least one fails.</p>
 * <p>
 * In case of multiple batch failures, which exception will be in the final {@link ListenableFuture} is
 * non-deterministic.
 * </p>
 * @param acknowledgeablePubsubMessages messages, could be from different subscriptions.
 * @param asyncOperation specific Pub/Sub operation to perform.
 * @return {@link ListenableFuture} indicating overall success or failure.
 */
private ListenableFuture<Void> doBatchedAsyncOperation(Collection<? extends AcknowledgeablePubsubMessage> acknowledgeablePubsubMessages, BiFunction<String, List<String>, ApiFuture<Empty>> asyncOperation) {
    Map<ProjectSubscriptionName, List<String>> groupedMessages = acknowledgeablePubsubMessages.stream().collect(Collectors.groupingBy(AcknowledgeablePubsubMessage::getProjectSubscriptionName, Collectors.mapping(AcknowledgeablePubsubMessage::getAckId, Collectors.toList())));
    replacedert.state(groupedMessages.keySet().stream().map(ProjectSubscriptionName::getProject).distinct().count() == 1, "The project id of all messages must match.");
    SettableListenableFuture<Void> settableListenableFuture = new SettableListenableFuture<>();
    int numExpectedFutures = groupedMessages.size();
    AtomicInteger numCompletedFutures = new AtomicInteger();
    groupedMessages.forEach((ProjectSubscriptionName psName, List<String> ackIds) -> {
        ApiFuture<Empty> ackApiFuture = asyncOperation.apply(psName.toString(), ackIds);
        ApiFutures.addCallback(ackApiFuture, new ApiFutureCallback<Empty>() {

            @Override
            public void onFailure(Throwable throwable) {
                processResult(throwable);
            }

            @Override
            public void onSuccess(Empty empty) {
                processResult(null);
            }

            private void processResult(Throwable throwable) {
                if (throwable != null) {
                    settableListenableFuture.setException(throwable);
                } else if (numCompletedFutures.incrementAndGet() == numExpectedFutures) {
                    settableListenableFuture.set(null);
                }
            }
        }, this.ackExecutor);
    });
    return settableListenableFuture;
}

/**
 * Unit tests for {@link DefaultTransportRequest}.
 *
 * @author Rossen Stoyanchev
 */
public clreplaced DefaultTransportRequestTests {

    private static final Jackson2SockJsMessageCodec CODEC = new Jackson2SockJsMessageCodec();

    private SettableListenableFuture<WebSocketSession> connectFuture;

    private ListenableFutureCallback<WebSocketSession> connectCallback;

    private TestTransport webSocketTransport;

    private TestTransport xhrTransport;

    @SuppressWarnings("unchecked")
    @BeforeEach
    public void setup() throws Exception {
        this.connectCallback = mock(ListenableFutureCallback.clreplaced);
        this.connectFuture = new SettableListenableFuture<>();
        this.connectFuture.addCallback(this.connectCallback);
        this.webSocketTransport = new TestTransport("WebSocketTestTransport");
        this.xhrTransport = new TestTransport("XhrTestTransport");
    }

    @Test
    public void connect() throws Exception {
        DefaultTransportRequest request = createTransportRequest(this.webSocketTransport, TransportType.WEBSOCKET);
        request.connect(null, this.connectFuture);
        WebSocketSession session = mock(WebSocketSession.clreplaced);
        this.webSocketTransport.getConnectCallback().onSuccess(session);
        replacedertThat(this.connectFuture.get()).isSameAs(session);
    }

    @Test
    public void fallbackAfterTransportError() throws Exception {
        DefaultTransportRequest request1 = createTransportRequest(this.webSocketTransport, TransportType.WEBSOCKET);
        DefaultTransportRequest request2 = createTransportRequest(this.xhrTransport, TransportType.XHR_STREAMING);
        request1.setFallbackRequest(request2);
        request1.connect(null, this.connectFuture);
        // Transport error => fallback
        this.webSocketTransport.getConnectCallback().onFailure(new IOException("Fake exception 1"));
        replacedertThat(this.connectFuture.isDone()).isFalse();
        replacedertThat(this.xhrTransport.invoked()).isTrue();
        // Transport error => no more fallback
        this.xhrTransport.getConnectCallback().onFailure(new IOException("Fake exception 2"));
        replacedertThat(this.connectFuture.isDone()).isTrue();
        replacedertThatExceptionOfType(ExecutionException.clreplaced).isThrownBy(this.connectFuture::get).withMessageContaining("Fake exception 2");
    }

    @Test
    public void fallbackAfterTimeout() throws Exception {
        TaskScheduler scheduler = mock(TaskScheduler.clreplaced);
        Runnable sessionCleanupTask = mock(Runnable.clreplaced);
        DefaultTransportRequest request1 = createTransportRequest(this.webSocketTransport, TransportType.WEBSOCKET);
        DefaultTransportRequest request2 = createTransportRequest(this.xhrTransport, TransportType.XHR_STREAMING);
        request1.setFallbackRequest(request2);
        request1.setTimeoutScheduler(scheduler);
        request1.addTimeoutTask(sessionCleanupTask);
        request1.connect(null, this.connectFuture);
        replacedertThat(this.webSocketTransport.invoked()).isTrue();
        replacedertThat(this.xhrTransport.invoked()).isFalse();
        // Get and invoke the scheduled timeout task
        ArgumentCaptor<Runnable> taskCaptor = ArgumentCaptor.forClreplaced(Runnable.clreplaced);
        verify(scheduler).schedule(taskCaptor.capture(), any(Date.clreplaced));
        verifyNoMoreInteractions(scheduler);
        taskCaptor.getValue().run();
        replacedertThat(this.xhrTransport.invoked()).isTrue();
        verify(sessionCleanupTask).run();
    }

    protected DefaultTransportRequest createTransportRequest(Transport transport, TransportType type) throws Exception {
        SockJsUrlInfo urlInfo = new SockJsUrlInfo(new URI("https://example.com"));
        return new DefaultTransportRequest(urlInfo, new HttpHeaders(), new HttpHeaders(), transport, type, CODEC);
    }
}

@Test
public void sendWithExecutionException() {
    this.session.afterConnected(this.connection);
    replacedertThat(this.session.isConnected()).isTrue();
    IllegalStateException exception = new IllegalStateException("simulated exception");
    SettableListenableFuture<Void> future = new SettableListenableFuture<>();
    future.setException(exception);
    given(this.connection.send(any())).willReturn(future);
    replacedertThatExceptionOfType(MessageDeliveryException.clreplaced).isThrownBy(() -> this.session.send("/topic/foo", "sample payload".getBytes(StandardCharsets.UTF_8))).withCause(exception);
}

/**
 * Unit tests for {@link DefaultTransportRequest}.
 *
 * @author Rossen Stoyanchev
 */
public clreplaced DefaultTransportRequestTests {

    private static final Jackson2SockJsMessageCodec CODEC = new Jackson2SockJsMessageCodec();

    private SettableListenableFuture<WebSocketSession> connectFuture;

    private ListenableFutureCallback<WebSocketSession> connectCallback;

    private TestTransport webSocketTransport;

    private TestTransport xhrTransport;

    @Rule
    public final ExpectedException thrown = ExpectedException.none();

    @SuppressWarnings("unchecked")
    @Before
    public void setup() throws Exception {
        this.connectCallback = mock(ListenableFutureCallback.clreplaced);
        this.connectFuture = new SettableListenableFuture<>();
        this.connectFuture.addCallback(this.connectCallback);
        this.webSocketTransport = new TestTransport("WebSocketTestTransport");
        this.xhrTransport = new TestTransport("XhrTestTransport");
    }

    @Test
    public void connect() throws Exception {
        DefaultTransportRequest request = createTransportRequest(this.webSocketTransport, TransportType.WEBSOCKET);
        request.connect(null, this.connectFuture);
        WebSocketSession session = mock(WebSocketSession.clreplaced);
        this.webSocketTransport.getConnectCallback().onSuccess(session);
        replacedertSame(session, this.connectFuture.get());
    }

    @Test
    public void fallbackAfterTransportError() throws Exception {
        DefaultTransportRequest request1 = createTransportRequest(this.webSocketTransport, TransportType.WEBSOCKET);
        DefaultTransportRequest request2 = createTransportRequest(this.xhrTransport, TransportType.XHR_STREAMING);
        request1.setFallbackRequest(request2);
        request1.connect(null, this.connectFuture);
        // Transport error => fallback
        this.webSocketTransport.getConnectCallback().onFailure(new IOException("Fake exception 1"));
        replacedertFalse(this.connectFuture.isDone());
        replacedertTrue(this.xhrTransport.invoked());
        // Transport error => no more fallback
        this.xhrTransport.getConnectCallback().onFailure(new IOException("Fake exception 2"));
        replacedertTrue(this.connectFuture.isDone());
        this.thrown.expect(ExecutionException.clreplaced);
        this.thrown.expectMessage("Fake exception 2");
        this.connectFuture.get();
    }

    @Test
    public void fallbackAfterTimeout() throws Exception {
        TaskScheduler scheduler = mock(TaskScheduler.clreplaced);
        Runnable sessionCleanupTask = mock(Runnable.clreplaced);
        DefaultTransportRequest request1 = createTransportRequest(this.webSocketTransport, TransportType.WEBSOCKET);
        DefaultTransportRequest request2 = createTransportRequest(this.xhrTransport, TransportType.XHR_STREAMING);
        request1.setFallbackRequest(request2);
        request1.setTimeoutScheduler(scheduler);
        request1.addTimeoutTask(sessionCleanupTask);
        request1.connect(null, this.connectFuture);
        replacedertTrue(this.webSocketTransport.invoked());
        replacedertFalse(this.xhrTransport.invoked());
        // Get and invoke the scheduled timeout task
        ArgumentCaptor<Runnable> taskCaptor = ArgumentCaptor.forClreplaced(Runnable.clreplaced);
        verify(scheduler).schedule(taskCaptor.capture(), any(Date.clreplaced));
        verifyNoMoreInteractions(scheduler);
        taskCaptor.getValue().run();
        replacedertTrue(this.xhrTransport.invoked());
        verify(sessionCleanupTask).run();
    }

    protected DefaultTransportRequest createTransportRequest(Transport transport, TransportType type) throws Exception {
        SockJsUrlInfo urlInfo = new SockJsUrlInfo(new URI("http://example.com"));
        return new DefaultTransportRequest(urlInfo, new HttpHeaders(), new HttpHeaders(), transport, type, CODEC);
    }
}

@Test
public void sendWithExecutionException() {
    this.session.afterConnected(this.connection);
    replacedertTrue(this.session.isConnected());
    IllegalStateException exception = new IllegalStateException("simulated exception");
    SettableListenableFuture<Void> future = new SettableListenableFuture<>();
    future.setException(exception);
    when(this.connection.send(any())).thenReturn(future);
    this.expected.expect(MessageDeliveryException.clreplaced);
    this.expected.expectCause(Matchers.sameInstance(exception));
    this.session.send("/topic/foo", "sample payload".getBytes(StandardCharsets.UTF_8));
    verifyNoMoreInteractions(this.connection);
}

@Test
public void sendWithExecutionException() throws Exception {
    this.session.afterConnected(this.connection);
    replacedertTrue(this.session.isConnected());
    IllegalStateException exception = new IllegalStateException("simulated exception");
    SettableListenableFuture<Void> future = new SettableListenableFuture<>();
    future.setException(exception);
    when(this.connection.send(any())).thenReturn(future);
    this.expected.expect(MessageDeliveryException.clreplaced);
    this.expected.expectCause(Matchers.sameInstance(exception));
    this.session.send("/topic/foo", "sample payload".getBytes(UTF_8));
    verifyNoMoreInteractions(this.connection);
}

/**
 */
clreplaced NettyResponseHandler extends SimpleChannelInboundHandler<HttpObject> {

    private final SettableListenableFuture<ClientHttpResponse> responseFuture;

    private final ChunkedInputStream<ByteBufHolder> in = new ChunkedInputStream<>(ByteBufHolderAdapter.INSTANCE);

    NettyResponseHandler(SettableListenableFuture<ClientHttpResponse> responseFuture) {
        this.responseFuture = responseFuture;
    }

    @Override
    protected void channelRead0(ChannelHandlerContext context, HttpObject response) throws Exception {
        if (response instanceof HttpResponse) {
            this.responseFuture.set(new NettyResponse(context, (HttpResponse) response, in));
        } else if (response instanceof HttpContent) {
            HttpContent cont = (HttpContent) response;
            in.add(cont);
            if (response instanceof LastHttpContent) {
                in.end();
            }
        } else {
            throw new RuntimeException("Unknown message: " + response);
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext context, Throwable cause) throws Exception {
        this.responseFuture.setException(cause);
    }
}

/**
 * place for log configuration, and also
 * workaround for access tom some spring beans from {@link com.codeabovelab.dm.common.log.AmqpAppender } which does
 * not managed by spring
 */
@Configuration
public clreplaced LogConfiguration implements ApplicationContextAware {

    static final SettableListenableFuture<ApplicationContext> DEFERRED_RESULT = new SettableListenableFuture<>();

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        DEFERRED_RESULT.set(applicationContext);
    }
}

/**
 * Unit tests for
 * {@link org.springframework.web.socket.sockjs.client.AbstractClientSockJsSession}.
 *
 * @author Rossen Stoyanchev
 */
public clreplaced ClientSockJsSessionTests {

    private static final Jackson2SockJsMessageCodec CODEC = new Jackson2SockJsMessageCodec();

    private TestClientSockJsSession session;

    private WebSocketHandler handler;

    private SettableListenableFuture<WebSocketSession> connectFuture;

    @Before
    public void setup() throws Exception {
        SockJsUrlInfo urlInfo = new SockJsUrlInfo(new URI("https://example.com"));
        Transport transport = mock(Transport.clreplaced);
        TransportRequest request = new DefaultTransportRequest(urlInfo, null, null, transport, TransportType.XHR, CODEC);
        this.handler = mock(WebSocketHandler.clreplaced);
        this.connectFuture = new SettableListenableFuture<>();
        this.session = new TestClientSockJsSession(request, this.handler, this.connectFuture);
    }

    @Test
    public void handleFrameOpen() throws Exception {
        replacedertThat(this.session.isOpen(), is(false));
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.isOpen(), is(true));
        replacedertTrue(this.connectFuture.isDone());
        replacedertThat(this.connectFuture.get(), sameInstance(this.session));
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameOpenWhenStatusNotNew() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.isOpen(), is(true));
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.disconnectStatus, equalTo(new CloseStatus(1006, "Server lost session")));
    }

    @Test
    public void handleFrameOpenWithWebSocketHandlerException() throws Exception {
        willThrow(new IllegalStateException("Fake error")).given(this.handler).afterConnectionEstablished(this.session);
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.isOpen(), is(true));
    }

    @Test
    public void handleFrameMessage() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.handleFrame(SockJsFrame.messageFrame(CODEC, "foo", "bar").getContent());
        verify(this.handler).afterConnectionEstablished(this.session);
        verify(this.handler).handleMessage(this.session, new TextMessage("foo"));
        verify(this.handler).handleMessage(this.session, new TextMessage("bar"));
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameMessageWhenNotOpen() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.close();
        reset(this.handler);
        this.session.handleFrame(SockJsFrame.messageFrame(CODEC, "foo", "bar").getContent());
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameMessageWithBadData() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.handleFrame("a['bad data");
        replacedertThat(this.session.isOpen(), equalTo(false));
        replacedertThat(this.session.disconnectStatus, equalTo(CloseStatus.BAD_DATA));
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameMessageWithWebSocketHandlerException() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        willThrow(new IllegalStateException("Fake error")).given(this.handler).handleMessage(this.session, new TextMessage("foo"));
        willThrow(new IllegalStateException("Fake error")).given(this.handler).handleMessage(this.session, new TextMessage("bar"));
        this.session.handleFrame(SockJsFrame.messageFrame(CODEC, "foo", "bar").getContent());
        replacedertThat(this.session.isOpen(), equalTo(true));
        verify(this.handler).afterConnectionEstablished(this.session);
        verify(this.handler).handleMessage(this.session, new TextMessage("foo"));
        verify(this.handler).handleMessage(this.session, new TextMessage("bar"));
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameClose() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.handleFrame(SockJsFrame.closeFrame(1007, "").getContent());
        replacedertThat(this.session.isOpen(), equalTo(false));
        replacedertThat(this.session.disconnectStatus, equalTo(new CloseStatus(1007, "")));
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleTransportError() throws Exception {
        final IllegalStateException ex = new IllegalStateException("Fake error");
        this.session.handleTransportError(ex);
        verify(this.handler).handleTransportError(this.session, ex);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void afterTransportClosed() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.afterTransportClosed(CloseStatus.SERVER_ERROR);
        replacedertThat(this.session.isOpen(), equalTo(false));
        verify(this.handler).afterConnectionEstablished(this.session);
        verify(this.handler).afterConnectionClosed(this.session, CloseStatus.SERVER_ERROR);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void close() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.close();
        replacedertThat(this.session.isOpen(), equalTo(false));
        replacedertThat(this.session.disconnectStatus, equalTo(CloseStatus.NORMAL));
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void closeWithStatus() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.close(new CloseStatus(3000, "reason"));
        replacedertThat(this.session.disconnectStatus, equalTo(new CloseStatus(3000, "reason")));
    }

    @Test
    public void closeWithNullStatus() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThatIllegalArgumentException().isThrownBy(() -> this.session.close(null)).withMessageContaining("Invalid close status");
    }

    @Test
    public void closeWithStatusOutOfRange() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThatIllegalArgumentException().isThrownBy(() -> this.session.close(new CloseStatus(2999, "reason"))).withMessageContaining("Invalid close status");
    }

    @Test
    public void timeoutTask() {
        this.session.getTimeoutTask().run();
        replacedertThat(this.session.disconnectStatus, equalTo(new CloseStatus(2007, "Transport timed out")));
    }

    @Test
    public void send() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.sendMessage(new TextMessage("foo"));
        replacedertThat(this.session.sentMessage, equalTo(new TextMessage("[\"foo\"]")));
    }

    private static clreplaced TestClientSockJsSession extends AbstractClientSockJsSession {

        private TextMessage sentMessage;

        private CloseStatus disconnectStatus;

        protected TestClientSockJsSession(TransportRequest request, WebSocketHandler handler, SettableListenableFuture<WebSocketSession> connectFuture) {
            super(request, handler, connectFuture);
        }

        @Override
        protected void sendInternal(TextMessage textMessage) throws IOException {
            this.sentMessage = textMessage;
        }

        @Override
        protected void disconnect(CloseStatus status) throws IOException {
            this.disconnectStatus = status;
        }

        @Override
        public InetSocketAddress getLocalAddress() {
            return null;
        }

        @Override
        public InetSocketAddress getRemoteAddress() {
            return null;
        }

        @Override
        public String getAcceptedProtocol() {
            return null;
        }

        @Override
        public void setTextMessageSizeLimit(int messageSizeLimit) {
        }

        @Override
        public int getTextMessageSizeLimit() {
            return 0;
        }

        @Override
        public void setBinaryMessageSizeLimit(int messageSizeLimit) {
        }

        @Override
        public int getBinaryMessageSizeLimit() {
            return 0;
        }

        @Override
        public List<WebSocketExtension> getExtensions() {
            return null;
        }
    }
}

@Test
public void listenableFutureWithError() throws Exception {
    SettableListenableFuture<String> future = new SettableListenableFuture<>();
    IllegalStateException ex = new IllegalStateException();
    testHandle(future, ListenableFuture.clreplaced, () -> future.setException(ex), ex);
}

@Before
public void setUp() {
    MockitoAnnotations.initMocks(this);
    this.sessionHandler = mock(StompSessionHandler.clreplaced);
    this.connectHeaders = new StompHeaders();
    this.session = new DefaultStompSession(this.sessionHandler, this.connectHeaders);
    this.session.setMessageConverter(new StringMessageConverter());
    SettableListenableFuture<Void> future = new SettableListenableFuture<>();
    future.set(null);
    given(this.connection.send(this.messageCaptor.capture())).willReturn(future);
}

@Override
public ListenableFuture<Void> shutdown() {
    if (this.stopping) {
        SettableListenableFuture<Void> future = new SettableListenableFuture<>();
        future.set(null);
        return future;
    }
    this.stopping = true;
    Mono<Void> result;
    if (this.channelGroup != null) {
        result = FutureMono.from(this.channelGroup.close());
        if (this.loopResources != null) {
            result = result.onErrorResume(ex -> Mono.empty()).then(this.loopResources.disposeLater());
        }
        if (this.poolResources != null) {
            result = result.onErrorResume(ex -> Mono.empty()).then(this.poolResources.disposeLater());
        }
        result = result.onErrorResume(ex -> Mono.empty()).then(stopScheduler());
    } else {
        result = stopScheduler();
    }
    return new MonoToListenableFutureAdapter<>(result);
}

@Before
public void setUp() {
    this.message = new GenericMessage<byte[]>("testPayload".getBytes(), new MapBuilder<String, Object>().put("key1", "value1").put("key2", "value2").build());
    SettableListenableFuture<String> future = new SettableListenableFuture<>();
    future.set("benfica");
    when(this.pubSubTemplate.publish(eq("testTopic"), eq("testPayload".getBytes()), isA(Map.clreplaced))).thenReturn(future);
    this.adapter = new PubSubMessageHandler(this.pubSubTemplate, "testTopic");
}

@Override
public ListenableFuture<List<PubsubMessage>> pullAndAckAsync(String subscription, Integer maxMessages, Boolean returnImmediately) {
    PullRequest pullRequest = this.subscriberFactory.createPullRequest(subscription, maxMessages, returnImmediately);
    final SettableListenableFuture<List<PubsubMessage>> settableFuture = new SettableListenableFuture<>();
    this.pullAsync(pullRequest).addCallback(ackableMessages -> {
        if (!ackableMessages.isEmpty()) {
            ack(ackableMessages);
        }
        List<PubsubMessage> messages = ackableMessages.stream().map(AcknowledgeablePubsubMessage::getPubsubMessage).collect(Collectors.toList());
        settableFuture.set(messages);
    }, settableFuture::setException);
    return settableFuture;
}

@Override
public ListenableFuture<PubsubMessage> pullNextAsync(String subscription) {
    final SettableListenableFuture<PubsubMessage> settableFuture = new SettableListenableFuture<>();
    this.pullAndAckAsync(subscription, 1, true).addCallback(messages -> {
        PubsubMessage message = messages.isEmpty() ? null : messages.get(0);
        settableFuture.set(message);
    }, settableFuture::setException);
    return settableFuture;
}

@Override
public <T> ListenableFuture<List<ConvertedAcknowledgeablePubsubMessage<T>>> pullAndConvertAsync(String subscription, Integer maxMessages, Boolean returnImmediately, Clreplaced<T> payloadType) {
    final SettableListenableFuture<List<ConvertedAcknowledgeablePubsubMessage<T>>> settableFuture = new SettableListenableFuture<>();
    this.pullAsync(subscription, maxMessages, returnImmediately).addCallback(ackableMessages -> settableFuture.set(this.toConvertedAcknowledgeablePubsubMessages(payloadType, ackableMessages)), settableFuture::setException);
    return settableFuture;
}

/**
 * Pulls messages asynchronously, on demand, using the pull request in argument.
 *
 * @param pullRequest pull request containing the subscription name
 * @return the ListenableFuture for the asynchronous execution, returning
 * the list of {@link AcknowledgeablePubsubMessage} containing the ack ID, subscription
 * and acknowledger
 */
private ListenableFuture<List<AcknowledgeablePubsubMessage>> pullAsync(PullRequest pullRequest) {
    replacedert.notNull(pullRequest, "The pull request can't be null.");
    ApiFuture<PullResponse> pullFuture = this.subscriberStub.pullCallable().futureCall(pullRequest);
    final SettableListenableFuture<List<AcknowledgeablePubsubMessage>> settableFuture = new SettableListenableFuture<>();
    ApiFutures.addCallback(pullFuture, new ApiFutureCallback<PullResponse>() {

        @Override
        public void onFailure(Throwable throwable) {
            settableFuture.setException(throwable);
        }

        @Override
        public void onSuccess(PullResponse pullResponse) {
            List<AcknowledgeablePubsubMessage> result = toAcknowledgeablePubsubMessageList(pullResponse.getReceivedMessagesList(), pullRequest.getSubscription());
            settableFuture.set(result);
        }
    }, asyncPullExecutor);
    return settableFuture;
}

/**
 * Unit tests for
 * {@link org.springframework.web.socket.sockjs.client.AbstractClientSockJsSession}.
 *
 * @author Rossen Stoyanchev
 */
public clreplaced ClientSockJsSessionTests {

    private static final Jackson2SockJsMessageCodec CODEC = new Jackson2SockJsMessageCodec();

    private TestClientSockJsSession session;

    private WebSocketHandler handler;

    private SettableListenableFuture<WebSocketSession> connectFuture;

    @BeforeEach
    public void setup() throws Exception {
        SockJsUrlInfo urlInfo = new SockJsUrlInfo(new URI("https://example.com"));
        Transport transport = mock(Transport.clreplaced);
        TransportRequest request = new DefaultTransportRequest(urlInfo, null, null, transport, TransportType.XHR, CODEC);
        this.handler = mock(WebSocketHandler.clreplaced);
        this.connectFuture = new SettableListenableFuture<>();
        this.session = new TestClientSockJsSession(request, this.handler, this.connectFuture);
    }

    @Test
    public void handleFrameOpen() throws Exception {
        replacedertThat(this.session.isOpen()).isFalse();
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.isOpen()).isTrue();
        replacedertThat(this.connectFuture.isDone()).isTrue();
        replacedertThat(this.connectFuture.get()).isSameAs(this.session);
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameOpenWhenStatusNotNew() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.isOpen()).isTrue();
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.disconnectStatus).isEqualTo(new CloseStatus(1006, "Server lost session"));
    }

    @Test
    public void handleFrameOpenWithWebSocketHandlerException() throws Exception {
        willThrow(new IllegalStateException("Fake error")).given(this.handler).afterConnectionEstablished(this.session);
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.isOpen()).isTrue();
    }

    @Test
    public void handleFrameMessage() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.handleFrame(SockJsFrame.messageFrame(CODEC, "foo", "bar").getContent());
        verify(this.handler).afterConnectionEstablished(this.session);
        verify(this.handler).handleMessage(this.session, new TextMessage("foo"));
        verify(this.handler).handleMessage(this.session, new TextMessage("bar"));
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameMessageWhenNotOpen() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.close();
        reset(this.handler);
        this.session.handleFrame(SockJsFrame.messageFrame(CODEC, "foo", "bar").getContent());
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameMessageWithBadData() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.handleFrame("a['bad data");
        replacedertThat(this.session.isOpen()).isEqualTo(false);
        replacedertThat(this.session.disconnectStatus).isEqualTo(CloseStatus.BAD_DATA);
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameMessageWithWebSocketHandlerException() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        willThrow(new IllegalStateException("Fake error")).given(this.handler).handleMessage(this.session, new TextMessage("foo"));
        willThrow(new IllegalStateException("Fake error")).given(this.handler).handleMessage(this.session, new TextMessage("bar"));
        this.session.handleFrame(SockJsFrame.messageFrame(CODEC, "foo", "bar").getContent());
        replacedertThat(this.session.isOpen()).isEqualTo(true);
        verify(this.handler).afterConnectionEstablished(this.session);
        verify(this.handler).handleMessage(this.session, new TextMessage("foo"));
        verify(this.handler).handleMessage(this.session, new TextMessage("bar"));
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameClose() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.handleFrame(SockJsFrame.closeFrame(1007, "").getContent());
        replacedertThat(this.session.isOpen()).isEqualTo(false);
        replacedertThat(this.session.disconnectStatus).isEqualTo(new CloseStatus(1007, ""));
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleTransportError() throws Exception {
        final IllegalStateException ex = new IllegalStateException("Fake error");
        this.session.handleTransportError(ex);
        verify(this.handler).handleTransportError(this.session, ex);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void afterTransportClosed() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.afterTransportClosed(CloseStatus.SERVER_ERROR);
        replacedertThat(this.session.isOpen()).isEqualTo(false);
        verify(this.handler).afterConnectionEstablished(this.session);
        verify(this.handler).afterConnectionClosed(this.session, CloseStatus.SERVER_ERROR);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void close() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.close();
        replacedertThat(this.session.isOpen()).isEqualTo(false);
        replacedertThat(this.session.disconnectStatus).isEqualTo(CloseStatus.NORMAL);
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void closeWithStatus() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.close(new CloseStatus(3000, "reason"));
        replacedertThat(this.session.disconnectStatus).isEqualTo(new CloseStatus(3000, "reason"));
    }

    @Test
    public void closeWithNullStatus() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThatIllegalArgumentException().isThrownBy(() -> this.session.close(null)).withMessageContaining("Invalid close status");
    }

    @Test
    public void closeWithStatusOutOfRange() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThatIllegalArgumentException().isThrownBy(() -> this.session.close(new CloseStatus(2999, "reason"))).withMessageContaining("Invalid close status");
    }

    @Test
    public void timeoutTask() {
        this.session.getTimeoutTask().run();
        replacedertThat(this.session.disconnectStatus).isEqualTo(new CloseStatus(2007, "Transport timed out"));
    }

    @Test
    public void send() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.sendMessage(new TextMessage("foo"));
        replacedertThat(this.session.sentMessage).isEqualTo(new TextMessage("[\"foo\"]"));
    }

    private static clreplaced TestClientSockJsSession extends AbstractClientSockJsSession {

        private TextMessage sentMessage;

        private CloseStatus disconnectStatus;

        protected TestClientSockJsSession(TransportRequest request, WebSocketHandler handler, SettableListenableFuture<WebSocketSession> connectFuture) {
            super(request, handler, connectFuture);
        }

        @Override
        protected void sendInternal(TextMessage textMessage) throws IOException {
            this.sentMessage = textMessage;
        }

        @Override
        protected void disconnect(CloseStatus status) throws IOException {
            this.disconnectStatus = status;
        }

        @Override
        public InetSocketAddress getLocalAddress() {
            return null;
        }

        @Override
        public InetSocketAddress getRemoteAddress() {
            return null;
        }

        @Override
        public String getAcceptedProtocol() {
            return null;
        }

        @Override
        public void setTextMessageSizeLimit(int messageSizeLimit) {
        }

        @Override
        public int getTextMessageSizeLimit() {
            return 0;
        }

        @Override
        public void setBinaryMessageSizeLimit(int messageSizeLimit) {
        }

        @Override
        public int getBinaryMessageSizeLimit() {
            return 0;
        }

        @Override
        public List<WebSocketExtension> getExtensions() {
            return null;
        }
    }
}

@BeforeEach
public void setUp() {
    this.connectHeaders = new StompHeaders();
    this.session = new DefaultStompSession(this.sessionHandler, this.connectHeaders);
    this.session.setMessageConverter(new CompositeMessageConverter(Arrays.asList(new StringMessageConverter(), new ByteArrayMessageConverter())));
    SettableListenableFuture<Void> future = new SettableListenableFuture<>();
    future.set(null);
    given(this.connection.send(this.messageCaptor.capture())).willReturn(future);
}

@Test
public void listenableFutureReply() {
    load(TestEventListener.clreplaced, ReplyEventListener.clreplaced);
    SettableListenableFuture<String> future = new SettableListenableFuture<>();
    future.set("dummy");
    AnotherTestEvent event = new AnotherTestEvent(this, future);
    ReplyEventListener replyEventListener = this.context.getBean(ReplyEventListener.clreplaced);
    TestEventListener listener = this.context.getBean(TestEventListener.clreplaced);
    this.eventCollector.replacedertNoEventReceived(listener);
    this.eventCollector.replacedertNoEventReceived(replyEventListener);
    this.context.publishEvent(event);
    this.eventCollector.replacedertEvent(replyEventListener, event);
    // reply
    this.eventCollector.replacedertEvent(listener, "dummy");
    this.eventCollector.replacedertTotalEventsCount(2);
}

/**
 * Unit tests for
 * {@link org.springframework.web.socket.sockjs.client.AbstractClientSockJsSession}.
 *
 * @author Rossen Stoyanchev
 */
public clreplaced ClientSockJsSessionTests {

    private static final Jackson2SockJsMessageCodec CODEC = new Jackson2SockJsMessageCodec();

    private TestClientSockJsSession session;

    private WebSocketHandler handler;

    private SettableListenableFuture<WebSocketSession> connectFuture;

    @Rule
    public final ExpectedException thrown = ExpectedException.none();

    @Before
    public void setup() throws Exception {
        SockJsUrlInfo urlInfo = new SockJsUrlInfo(new URI("http://example.com"));
        Transport transport = mock(Transport.clreplaced);
        TransportRequest request = new DefaultTransportRequest(urlInfo, null, null, transport, TransportType.XHR, CODEC);
        this.handler = mock(WebSocketHandler.clreplaced);
        this.connectFuture = new SettableListenableFuture<>();
        this.session = new TestClientSockJsSession(request, this.handler, this.connectFuture);
    }

    @Test
    public void handleFrameOpen() throws Exception {
        replacedertThat(this.session.isOpen(), is(false));
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.isOpen(), is(true));
        replacedertTrue(this.connectFuture.isDone());
        replacedertThat(this.connectFuture.get(), sameInstance(this.session));
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameOpenWhenStatusNotNew() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.isOpen(), is(true));
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.disconnectStatus, equalTo(new CloseStatus(1006, "Server lost session")));
    }

    @Test
    public void handleFrameOpenWithWebSocketHandlerException() throws Exception {
        willThrow(new IllegalStateException("Fake error")).given(this.handler).afterConnectionEstablished(this.session);
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        replacedertThat(this.session.isOpen(), is(true));
    }

    @Test
    public void handleFrameMessage() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.handleFrame(SockJsFrame.messageFrame(CODEC, "foo", "bar").getContent());
        verify(this.handler).afterConnectionEstablished(this.session);
        verify(this.handler).handleMessage(this.session, new TextMessage("foo"));
        verify(this.handler).handleMessage(this.session, new TextMessage("bar"));
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameMessageWhenNotOpen() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.close();
        reset(this.handler);
        this.session.handleFrame(SockJsFrame.messageFrame(CODEC, "foo", "bar").getContent());
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameMessageWithBadData() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.handleFrame("a['bad data");
        replacedertThat(this.session.isOpen(), equalTo(false));
        replacedertThat(this.session.disconnectStatus, equalTo(CloseStatus.BAD_DATA));
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameMessageWithWebSocketHandlerException() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        willThrow(new IllegalStateException("Fake error")).given(this.handler).handleMessage(this.session, new TextMessage("foo"));
        willThrow(new IllegalStateException("Fake error")).given(this.handler).handleMessage(this.session, new TextMessage("bar"));
        this.session.handleFrame(SockJsFrame.messageFrame(CODEC, "foo", "bar").getContent());
        replacedertThat(this.session.isOpen(), equalTo(true));
        verify(this.handler).afterConnectionEstablished(this.session);
        verify(this.handler).handleMessage(this.session, new TextMessage("foo"));
        verify(this.handler).handleMessage(this.session, new TextMessage("bar"));
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleFrameClose() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.handleFrame(SockJsFrame.closeFrame(1007, "").getContent());
        replacedertThat(this.session.isOpen(), equalTo(false));
        replacedertThat(this.session.disconnectStatus, equalTo(new CloseStatus(1007, "")));
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void handleTransportError() throws Exception {
        final IllegalStateException ex = new IllegalStateException("Fake error");
        this.session.handleTransportError(ex);
        verify(this.handler).handleTransportError(this.session, ex);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void afterTransportClosed() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.afterTransportClosed(CloseStatus.SERVER_ERROR);
        replacedertThat(this.session.isOpen(), equalTo(false));
        verify(this.handler).afterConnectionEstablished(this.session);
        verify(this.handler).afterConnectionClosed(this.session, CloseStatus.SERVER_ERROR);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void close() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.close();
        replacedertThat(this.session.isOpen(), equalTo(false));
        replacedertThat(this.session.disconnectStatus, equalTo(CloseStatus.NORMAL));
        verify(this.handler).afterConnectionEstablished(this.session);
        verifyNoMoreInteractions(this.handler);
    }

    @Test
    public void closeWithStatus() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.close(new CloseStatus(3000, "reason"));
        replacedertThat(this.session.disconnectStatus, equalTo(new CloseStatus(3000, "reason")));
    }

    @Test
    public void closeWithNullStatus() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.thrown.expect(IllegalArgumentException.clreplaced);
        this.thrown.expectMessage("Invalid close status");
        this.session.close(null);
    }

    @Test
    public void closeWithStatusOutOfRange() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.thrown.expect(IllegalArgumentException.clreplaced);
        this.thrown.expectMessage("Invalid close status");
        this.session.close(new CloseStatus(2999, "reason"));
    }

    @Test
    public void timeoutTask() {
        this.session.getTimeoutTask().run();
        replacedertThat(this.session.disconnectStatus, equalTo(new CloseStatus(2007, "Transport timed out")));
    }

    @Test
    public void send() throws Exception {
        this.session.handleFrame(SockJsFrame.openFrame().getContent());
        this.session.sendMessage(new TextMessage("foo"));
        replacedertThat(this.session.sentMessage, equalTo(new TextMessage("[\"foo\"]")));
    }

    private static clreplaced TestClientSockJsSession extends AbstractClientSockJsSession {

        private TextMessage sentMessage;

        private CloseStatus disconnectStatus;

        protected TestClientSockJsSession(TransportRequest request, WebSocketHandler handler, SettableListenableFuture<WebSocketSession> connectFuture) {
            super(request, handler, connectFuture);
        }

        @Override
        protected void sendInternal(TextMessage textMessage) throws IOException {
            this.sentMessage = textMessage;
        }

        @Override
        protected void disconnect(CloseStatus status) throws IOException {
            this.disconnectStatus = status;
        }

        @Override
        public InetSocketAddress getLocalAddress() {
            return null;
        }

        @Override
        public InetSocketAddress getRemoteAddress() {
            return null;
        }

        @Override
        public String getAcceptedProtocol() {
            return null;
        }

        @Override
        public void setTextMessageSizeLimit(int messageSizeLimit) {
        }

        @Override
        public int getTextMessageSizeLimit() {
            return 0;
        }

        @Override
        public void setBinaryMessageSizeLimit(int messageSizeLimit) {
        }

        @Override
        public int getBinaryMessageSizeLimit() {
            return 0;
        }

        @Override
        public List<WebSocketExtension> getExtensions() {
            return null;
        }
    }
}

@Before
public void setUp() {
    MockitoAnnotations.initMocks(this);
    this.sessionHandler = mock(StompSessionHandler.clreplaced);
    this.connectHeaders = new StompHeaders();
    this.session = new DefaultStompSession(this.sessionHandler, this.connectHeaders);
    this.session.setMessageConverter(new StringMessageConverter());
    SettableListenableFuture<Void> future = new SettableListenableFuture<>();
    future.set(null);
    when(this.connection.send(this.messageCaptor.capture())).thenReturn(future);
}

@Before
public void setUp() throws Exception {
    MockitoAnnotations.initMocks(this);
    this.sessionHandler = mock(StompSessionHandler.clreplaced);
    this.connectHeaders = new StompHeaders();
    this.session = new DefaultStompSession(this.sessionHandler, this.connectHeaders);
    this.session.setMessageConverter(new StringMessageConverter());
    SettableListenableFuture<Void> future = new SettableListenableFuture<>();
    future.set(null);
    when(this.connection.send(this.messageCaptor.capture())).thenReturn(future);
}

public static <T> ListenableFuture<T> settable(T toSet) {
    SettableListenableFuture<T> settable = new SettableListenableFuture<>();
    settable.set(toSet);
    return settable;
}

@Before
public void setUp() {
    this.message = new GenericMessage<byte[]>("testPayload".getBytes(), new MapBuilder<String, Object>().put("key1", "value1").put("key2", "value2").build());
    SettableListenableFuture<String> future = new SettableListenableFuture<>();
    future.set("benfica");
    when(this.pubSubTemplate.publish(eq("testTopic"), eq("testPayload".getBytes()), anyMap())).thenReturn(future);
    this.adapter = new PubSubMessageHandler(this.pubSubTemplate, "testTopic");
}

protected ListenableFuture<ClientHttpResponse> executeInternal(final HttpHeaders headers) {
    final SettableListenableFuture<ClientHttpResponse> responseFuture = new SettableListenableFuture<>();
    ChannelFutureListener connectionListener = future -> {
        if (future.isSuccess()) {
            Channel channel = future.channel();
            channel.pipeline().addLast(new NettyResponseHandler(responseFuture));
            FullHttpRequest nettyRequest = createFullHttpRequest(headers);
            channel.writeAndFlush(nettyRequest);
        } else {
            responseFuture.setException(future.cause());
        }
    };
    this.bootstrap.connect(this.uri.getHost(), getPort(this.uri)).addListener(connectionListener);
    return responseFuture;
}

/**
 * Base clreplaced for SockJS client implementations of {@link WebSocketSession}.
 * Provides processing of incoming SockJS message frames and delegates lifecycle
 * events and messages to the (application) {@link WebSocketHandler}.
 * Sub-clreplacedes implement actual send as well as disconnect logic.
 *
 * @author Rossen Stoyanchev
 * @author Juergen Hoeller
 * @since 4.1
 */
public abstract clreplaced AbstractClientSockJsSession implements WebSocketSession {

    protected final Log logger = LogFactory.getLog(getClreplaced());

    private final TransportRequest request;

    private final WebSocketHandler webSocketHandler;

    private final SettableListenableFuture<WebSocketSession> connectFuture;

    private final Map<String, Object> attributes = new ConcurrentHashMap<>();

    @Nullable
    private volatile State state = State.NEW;

    @Nullable
    private volatile CloseStatus closeStatus;

    protected AbstractClientSockJsSession(TransportRequest request, WebSocketHandler handler, SettableListenableFuture<WebSocketSession> connectFuture) {
        replacedert.notNull(request, "'request' is required");
        replacedert.notNull(handler, "'handler' is required");
        replacedert.notNull(connectFuture, "'connectFuture' is required");
        this.request = request;
        this.webSocketHandler = handler;
        this.connectFuture = connectFuture;
    }

    @Override
    public String getId() {
        return this.request.getSockJsUrlInfo().getSessionId();
    }

    @Override
    public URI getUri() {
        return this.request.getSockJsUrlInfo().getSockJsUrl();
    }

    @Override
    public HttpHeaders getHandshakeHeaders() {
        return this.request.getHandshakeHeaders();
    }

    @Override
    public Map<String, Object> getAttributes() {
        return this.attributes;
    }

    @Override
    public Principal getPrincipal() {
        return this.request.getUser();
    }

    public SockJsMessageCodec getMessageCodec() {
        return this.request.getMessageCodec();
    }

    public WebSocketHandler getWebSocketHandler() {
        return this.webSocketHandler;
    }

    /**
     * Return a timeout cleanup task to invoke if the SockJS sessions is not
     * fully established within the retransmission timeout period calculated in
     * {@code SockJsRequest} based on the duration of the initial SockJS "Info"
     * request.
     */
    Runnable getTimeoutTask() {
        return new Runnable() {

            @Override
            public void run() {
                try {
                    closeInternal(new CloseStatus(2007, "Transport timed out"));
                } catch (Throwable ex) {
                    if (logger.isWarnEnabled()) {
                        logger.warn("Failed to close " + this + " after transport timeout", ex);
                    }
                }
            }
        };
    }

    @Override
    public boolean isOpen() {
        return (this.state == State.OPEN);
    }

    public boolean isDisconnected() {
        return (this.state == State.CLOSING || this.state == State.CLOSED);
    }

    @Override
    public final void sendMessage(WebSocketMessage<?> message) throws IOException {
        if (!(message instanceof TextMessage)) {
            throw new IllegalArgumentException(this + " supports text messages only.");
        }
        if (this.state != State.OPEN) {
            throw new IllegalStateException(this + " is not open: current state " + this.state);
        }
        String payload = ((TextMessage) message).getPayload();
        payload = getMessageCodec().encode(payload);
        // the client-side doesn't need message framing (letter "a")
        payload = payload.substring(1);
        TextMessage messageToSend = new TextMessage(payload);
        if (logger.isTraceEnabled()) {
            logger.trace("Sending message " + messageToSend + " in " + this);
        }
        sendInternal(messageToSend);
    }

    protected abstract void sendInternal(TextMessage textMessage) throws IOException;

    @Override
    public final void close() throws IOException {
        close(CloseStatus.NORMAL);
    }

    @Override
    public final void close(CloseStatus status) throws IOException {
        if (!isUserSetStatus(status)) {
            throw new IllegalArgumentException("Invalid close status: " + status);
        }
        if (logger.isDebugEnabled()) {
            logger.debug("Closing session with " + status + " in " + this);
        }
        closeInternal(status);
    }

    private boolean isUserSetStatus(@Nullable CloseStatus status) {
        return (status != null && (status.getCode() == 1000 || (status.getCode() >= 3000 && status.getCode() <= 4999)));
    }

    private void silentClose(CloseStatus status) {
        try {
            closeInternal(status);
        } catch (Throwable ex) {
            if (logger.isWarnEnabled()) {
                logger.warn("Failed to close " + this, ex);
            }
        }
    }

    protected void closeInternal(CloseStatus status) throws IOException {
        if (this.state == null) {
            logger.warn("Ignoring close since connect() was never invoked");
            return;
        }
        if (isDisconnected()) {
            if (logger.isDebugEnabled()) {
                logger.debug("Ignoring close (already closing or closed): current state " + this.state);
            }
            return;
        }
        this.state = State.CLOSING;
        this.closeStatus = status;
        disconnect(status);
    }

    protected abstract void disconnect(CloseStatus status) throws IOException;

    public void handleFrame(String payload) {
        SockJsFrame frame = new SockJsFrame(payload);
        switch(frame.getType()) {
            case OPEN:
                handleOpenFrame();
                break;
            case HEARTBEAT:
                if (logger.isTraceEnabled()) {
                    logger.trace("Received heartbeat in " + this);
                }
                break;
            case MESSAGE:
                handleMessageFrame(frame);
                break;
            case CLOSE:
                handleCloseFrame(frame);
        }
    }

    private void handleOpenFrame() {
        if (logger.isDebugEnabled()) {
            logger.debug("Processing SockJS open frame in " + this);
        }
        if (this.state == State.NEW) {
            this.state = State.OPEN;
            try {
                this.webSocketHandler.afterConnectionEstablished(this);
                this.connectFuture.set(this);
            } catch (Throwable ex) {
                if (logger.isErrorEnabled()) {
                    logger.error("WebSocketHandler.afterConnectionEstablished threw exception in " + this, ex);
                }
            }
        } else {
            if (logger.isDebugEnabled()) {
                logger.debug("Open frame received in " + getId() + " but we're not connecting (current state " + this.state + "). The server might have been restarted and lost track of the session.");
            }
            silentClose(new CloseStatus(1006, "Server lost session"));
        }
    }

    private void handleMessageFrame(SockJsFrame frame) {
        if (!isOpen()) {
            if (logger.isErrorEnabled()) {
                logger.error("Ignoring received message due to state " + this.state + " in " + this);
            }
            return;
        }
        String[] messages = null;
        String frameData = frame.getFrameData();
        if (frameData != null) {
            try {
                messages = getMessageCodec().decode(frameData);
            } catch (IOException ex) {
                if (logger.isErrorEnabled()) {
                    logger.error("Failed to decode data for SockJS \"message\" frame: " + frame + " in " + this, ex);
                }
                silentClose(CloseStatus.BAD_DATA);
                return;
            }
        }
        if (messages == null) {
            return;
        }
        if (logger.isTraceEnabled()) {
            logger.trace("Processing SockJS message frame " + frame.getContent() + " in " + this);
        }
        for (String message : messages) {
            if (isOpen()) {
                try {
                    this.webSocketHandler.handleMessage(this, new TextMessage(message));
                } catch (Throwable ex) {
                    logger.error("WebSocketHandler.handleMessage threw an exception on " + frame + " in " + this, ex);
                }
            }
        }
    }

    private void handleCloseFrame(SockJsFrame frame) {
        CloseStatus closeStatus = CloseStatus.NO_STATUS_CODE;
        try {
            String frameData = frame.getFrameData();
            if (frameData != null) {
                String[] data = getMessageCodec().decode(frameData);
                if (data != null && data.length == 2) {
                    closeStatus = new CloseStatus(Integer.valueOf(data[0]), data[1]);
                }
                if (logger.isDebugEnabled()) {
                    logger.debug("Processing SockJS close frame with " + closeStatus + " in " + this);
                }
            }
        } catch (IOException ex) {
            if (logger.isErrorEnabled()) {
                logger.error("Failed to decode data for " + frame + " in " + this, ex);
            }
        }
        silentClose(closeStatus);
    }

    public void handleTransportError(Throwable error) {
        try {
            if (logger.isErrorEnabled()) {
                logger.error("Transport error in " + this, error);
            }
            this.webSocketHandler.handleTransportError(this, error);
        } catch (Throwable ex) {
            logger.error("WebSocketHandler.handleTransportError threw an exception", ex);
        }
    }

    public void afterTransportClosed(@Nullable CloseStatus closeStatus) {
        CloseStatus cs = this.closeStatus;
        if (cs == null) {
            cs = closeStatus;
            this.closeStatus = closeStatus;
        }
        replacedert.state(cs != null, "CloseStatus not available");
        if (logger.isDebugEnabled()) {
            logger.debug("Transport closed with " + cs + " in " + this);
        }
        this.state = State.CLOSED;
        try {
            this.webSocketHandler.afterConnectionClosed(this, cs);
        } catch (Throwable ex) {
            logger.error("WebSocketHandler.afterConnectionClosed threw an exception", ex);
        }
    }

    @Override
    public String toString() {
        return getClreplaced().getSimpleName() + "[id='" + getId() + ", url=" + getUri() + "]";
    }

    private enum State {

        NEW, OPEN, CLOSING, CLOSED
    }
}

@Test
public void listenableFuture() throws Exception {
    SettableListenableFuture<String> future = new SettableListenableFuture<>();
    testHandle(future, ListenableFuture.clreplaced, () -> future.set("foo"), "foo");
}

@Override
protected ListenableFuture<ClientHttpResponse> executeInternal(final HttpHeaders headers) throws IOException {
    final SettableListenableFuture<ClientHttpResponse> responseFuture = new SettableListenableFuture<>();
    ChannelFutureListener connectionListener = future -> {
        if (future.isSuccess()) {
            Channel channel = future.channel();
            channel.pipeline().addLast(new RequestExecuteHandler(responseFuture));
            FullHttpRequest nettyRequest = createFullHttpRequest(headers);
            channel.writeAndFlush(nettyRequest);
        } else {
            responseFuture.setException(future.cause());
        }
    };
    this.bootstrap.connect(this.uri.getHost(), getPort(this.uri)).addListener(connectionListener);
    return responseFuture;
}

@Override
public ListenableFuture<ClientHttpResponse> executeAsync() throws IOException {
    SettableListenableFuture<ClientHttpResponse> future = new SettableListenableFuture<>();
    future.set(execute());
    return future;
}

/**
 * Base clreplaced for SockJS client implementations of {@link WebSocketSession}.
 * Provides processing of incoming SockJS message frames and delegates lifecycle
 * events and messages to the (application) {@link WebSocketHandler}.
 * Sub-clreplacedes implement actual send as well as disconnect logic.
 *
 * @author Rossen Stoyanchev
 * @author Juergen Hoeller
 * @since 4.1
 */
public abstract clreplaced AbstractClientSockJsSession implements WebSocketSession {

    protected final Log logger = LogFactory.getLog(getClreplaced());

    private final TransportRequest request;

    private final WebSocketHandler webSocketHandler;

    private final SettableListenableFuture<WebSocketSession> connectFuture;

    private final Map<String, Object> attributes = new ConcurrentHashMap<>();

    @Nullable
    private volatile State state = State.NEW;

    @Nullable
    private volatile CloseStatus closeStatus;

    protected AbstractClientSockJsSession(TransportRequest request, WebSocketHandler handler, SettableListenableFuture<WebSocketSession> connectFuture) {
        replacedert.notNull(request, "'request' is required");
        replacedert.notNull(handler, "'handler' is required");
        replacedert.notNull(connectFuture, "'connectFuture' is required");
        this.request = request;
        this.webSocketHandler = handler;
        this.connectFuture = connectFuture;
    }

    @Override
    public String getId() {
        return this.request.getSockJsUrlInfo().getSessionId();
    }

    @Override
    public URI getUri() {
        return this.request.getSockJsUrlInfo().getSockJsUrl();
    }

    @Override
    public HttpHeaders getHandshakeHeaders() {
        return this.request.getHandshakeHeaders();
    }

    @Override
    public Map<String, Object> getAttributes() {
        return this.attributes;
    }

    @Override
    public Principal getPrincipal() {
        return this.request.getUser();
    }

    public SockJsMessageCodec getMessageCodec() {
        return this.request.getMessageCodec();
    }

    public WebSocketHandler getWebSocketHandler() {
        return this.webSocketHandler;
    }

    /**
     * Return a timeout cleanup task to invoke if the SockJS sessions is not
     * fully established within the retransmission timeout period calculated in
     * {@code SockJsRequest} based on the duration of the initial SockJS "Info"
     * request.
     */
    Runnable getTimeoutTask() {
        return new Runnable() {

            @Override
            public void run() {
                try {
                    closeInternal(new CloseStatus(2007, "Transport timed out"));
                } catch (Throwable ex) {
                    if (logger.isWarnEnabled()) {
                        logger.warn("Failed to close " + this + " after transport timeout", ex);
                    }
                }
            }
        };
    }

    @Override
    public boolean isOpen() {
        return (this.state == State.OPEN);
    }

    public boolean isDisconnected() {
        return (this.state == State.CLOSING || this.state == State.CLOSED);
    }

    @Override
    public final void sendMessage(WebSocketMessage<?> message) throws IOException {
        if (!(message instanceof TextMessage)) {
            throw new IllegalArgumentException(this + " supports text messages only.");
        }
        if (this.state != State.OPEN) {
            throw new IllegalStateException(this + " is not open: current state " + this.state);
        }
        String payload = ((TextMessage) message).getPayload();
        payload = getMessageCodec().encode(payload);
        // the client-side doesn't need message framing (letter "a")
        payload = payload.substring(1);
        TextMessage messageToSend = new TextMessage(payload);
        if (logger.isTraceEnabled()) {
            logger.trace("Sending message " + messageToSend + " in " + this);
        }
        sendInternal(messageToSend);
    }

    protected abstract void sendInternal(TextMessage textMessage) throws IOException;

    @Override
    public final void close() throws IOException {
        close(CloseStatus.NORMAL);
    }

    @Override
    public final void close(CloseStatus status) throws IOException {
        if (!isUserSetStatus(status)) {
            throw new IllegalArgumentException("Invalid close status: " + status);
        }
        if (logger.isDebugEnabled()) {
            logger.debug("Closing session with " + status + " in " + this);
        }
        closeInternal(status);
    }

    private boolean isUserSetStatus(@Nullable CloseStatus status) {
        return (status != null && (status.getCode() == 1000 || (status.getCode() >= 3000 && status.getCode() <= 4999)));
    }

    private void silentClose(CloseStatus status) {
        try {
            closeInternal(status);
        } catch (Throwable ex) {
            if (logger.isWarnEnabled()) {
                logger.warn("Failed to close " + this, ex);
            }
        }
    }

    protected void closeInternal(CloseStatus status) throws IOException {
        if (this.state == null) {
            logger.warn("Ignoring close since connect() was never invoked");
            return;
        }
        if (isDisconnected()) {
            if (logger.isDebugEnabled()) {
                logger.debug("Ignoring close (already closing or closed): current state " + this.state);
            }
            return;
        }
        this.state = State.CLOSING;
        this.closeStatus = status;
        disconnect(status);
    }

    protected abstract void disconnect(CloseStatus status) throws IOException;

    public void handleFrame(String payload) {
        SockJsFrame frame = new SockJsFrame(payload);
        switch(frame.getType()) {
            case OPEN:
                handleOpenFrame();
                break;
            case HEARTBEAT:
                if (logger.isTraceEnabled()) {
                    logger.trace("Received heartbeat in " + this);
                }
                break;
            case MESSAGE:
                handleMessageFrame(frame);
                break;
            case CLOSE:
                handleCloseFrame(frame);
        }
    }

    private void handleOpenFrame() {
        if (logger.isDebugEnabled()) {
            logger.debug("Processing SockJS open frame in " + this);
        }
        if (this.state == State.NEW) {
            this.state = State.OPEN;
            try {
                this.webSocketHandler.afterConnectionEstablished(this);
                this.connectFuture.set(this);
            } catch (Exception ex) {
                if (logger.isErrorEnabled()) {
                    logger.error("WebSocketHandler.afterConnectionEstablished threw exception in " + this, ex);
                }
            }
        } else {
            if (logger.isDebugEnabled()) {
                logger.debug("Open frame received in " + getId() + " but we're not connecting (current state " + this.state + "). The server might have been restarted and lost track of the session.");
            }
            silentClose(new CloseStatus(1006, "Server lost session"));
        }
    }

    private void handleMessageFrame(SockJsFrame frame) {
        if (!isOpen()) {
            if (logger.isErrorEnabled()) {
                logger.error("Ignoring received message due to state " + this.state + " in " + this);
            }
            return;
        }
        String[] messages = null;
        String frameData = frame.getFrameData();
        if (frameData != null) {
            try {
                messages = getMessageCodec().decode(frameData);
            } catch (IOException ex) {
                if (logger.isErrorEnabled()) {
                    logger.error("Failed to decode data for SockJS \"message\" frame: " + frame + " in " + this, ex);
                }
                silentClose(CloseStatus.BAD_DATA);
                return;
            }
        }
        if (messages == null) {
            return;
        }
        if (logger.isTraceEnabled()) {
            logger.trace("Processing SockJS message frame " + frame.getContent() + " in " + this);
        }
        for (String message : messages) {
            if (isOpen()) {
                try {
                    this.webSocketHandler.handleMessage(this, new TextMessage(message));
                } catch (Exception ex) {
                    logger.error("WebSocketHandler.handleMessage threw an exception on " + frame + " in " + this, ex);
                }
            }
        }
    }

    private void handleCloseFrame(SockJsFrame frame) {
        CloseStatus closeStatus = CloseStatus.NO_STATUS_CODE;
        try {
            String frameData = frame.getFrameData();
            if (frameData != null) {
                String[] data = getMessageCodec().decode(frameData);
                if (data != null && data.length == 2) {
                    closeStatus = new CloseStatus(Integer.parseInt(data[0]), data[1]);
                }
                if (logger.isDebugEnabled()) {
                    logger.debug("Processing SockJS close frame with " + closeStatus + " in " + this);
                }
            }
        } catch (IOException ex) {
            if (logger.isErrorEnabled()) {
                logger.error("Failed to decode data for " + frame + " in " + this, ex);
            }
        }
        silentClose(closeStatus);
    }

    public void handleTransportError(Throwable error) {
        try {
            if (logger.isErrorEnabled()) {
                logger.error("Transport error in " + this, error);
            }
            this.webSocketHandler.handleTransportError(this, error);
        } catch (Throwable ex) {
            logger.error("WebSocketHandler.handleTransportError threw an exception", ex);
        }
    }

    public void afterTransportClosed(@Nullable CloseStatus closeStatus) {
        CloseStatus cs = this.closeStatus;
        if (cs == null) {
            cs = closeStatus;
            this.closeStatus = closeStatus;
        }
        replacedert.state(cs != null, "CloseStatus not available");
        if (logger.isDebugEnabled()) {
            logger.debug("Transport closed with " + cs + " in " + this);
        }
        this.state = State.CLOSED;
        try {
            this.webSocketHandler.afterConnectionClosed(this, cs);
        } catch (Throwable ex) {
            logger.error("WebSocketHandler.afterConnectionClosed threw an exception", ex);
        }
    }

    @Override
    public String toString() {
        return getClreplaced().getSimpleName() + "[id='" + getId() + ", url=" + getUri() + "]";
    }

    private enum State {

        NEW, OPEN, CLOSING, CLOSED
    }
}

@Test
public void testSendQuestionnaireDataCallbackExceptionConsumed() {
    SettableListenableFuture<SendResult<String, RegisterIntegrationDTO>> resultFuture = new SettableListenableFuture<>();
    resultFuture.setException(new RuntimeException());
    when(kafkaTemplate.send(TOPIC_NAME, ENCRYPTED_MESSAGE)).thenReturn(resultFuture);
    registerIntegrationService.sendQuestionnaireData(questionnaireDoreplacedent, USER_GUID);
    verify(kafkaTemplate, times(1)).send(TOPIC_NAME, ENCRYPTED_MESSAGE);
}

@RunWith(MockitoJUnitRunner.clreplaced)
public clreplaced ResponseFutureAdapterTest {

    @Mock
    private ResponseEnreplacedy<String> responseEnreplacedy;

    @Mock
    private ListenableFutureCallback<String> callback;

    private SettableListenableFuture<ResponseEnreplacedy<String>> wrappedFuture = new SettableListenableFuture<>();

    @Before
    public void setUp() {
        ResponseFutureAdapter<String> responseFutureWrapper = new ResponseFutureAdapter<>(wrappedFuture);
        responseFutureWrapper.addCallback(callback);
    }

    @Test
    public void testOnSuccess() {
        String responseString = "SOMESTRING";
        when(responseEnreplacedy.getBody()).thenReturn(responseString);
        wrappedFuture.set(responseEnreplacedy);
        verify(callback).onSuccess(eq(responseString));
        verify(responseEnreplacedy).getBody();
    }

    @Test
    public void testOFailure() {
        Exception exception = new Exception();
        wrappedFuture.setException(exception);
        verify(callback).onFailure(eq(exception));
    }
}

@RunWith(MockitoJUnitRunner.clreplaced)
public clreplaced OptionalTypeFutureAdapterTest {

    @Mock
    private ListenableFutureCallback<Optional<String>> callback;

    private SettableListenableFuture<ResponseEnreplacedy<Optional<String>>> wrappedFuture = new SettableListenableFuture<>();

    @Before
    public void setUp() {
        OptionalTypeFutureAdapter<String> listenableFutureWrapper = new OptionalTypeFutureAdapter<>(wrappedFuture);
        listenableFutureWrapper.addCallback(callback);
    }

    @Test
    public void testOFailure() {
        Exception exception = new Exception();
        wrappedFuture.setException(exception);
        verify(callback).onFailure(eq(exception));
    }

    @Test
    public void testNotFoundExceptionTranslatedToEmptyOptional() {
        HttpClientErrorException exception = new HttpClientErrorException(HttpStatus.NOT_FOUND);
        wrappedFuture.setException(exception);
        verify(callback).onSuccess(eq(Optional.empty()));
    }

    @Test
    public void testOtherClientErrorExceptionsNotTranslatedToEmptyOptional() {
        HttpClientErrorException exception = new HttpClientErrorException(HttpStatus.INTERNAL_SERVER_ERROR);
        wrappedFuture.setException(exception);
        verify(callback).onFailure(eq(exception));
    }
}

/**
 * Base clreplaced for SockJS client implementations of {@link WebSocketSession}.
 * Provides processing of incoming SockJS message frames and delegates lifecycle
 * events and messages to the (application) {@link WebSocketHandler}.
 * Sub-clreplacedes implement actual send as well as disconnect logic.
 *
 * @author Rossen Stoyanchev
 * @since 4.1
 */
public abstract clreplaced AbstractClientSockJsSession implements WebSocketSession {

    protected final Log logger = LogFactory.getLog(getClreplaced());

    private final TransportRequest request;

    private final WebSocketHandler webSocketHandler;

    private final SettableListenableFuture<WebSocketSession> connectFuture;

    private final Map<String, Object> attributes = new ConcurrentHashMap<String, Object>();

    private volatile State state = State.NEW;

    private volatile CloseStatus closeStatus;

    protected AbstractClientSockJsSession(TransportRequest request, WebSocketHandler handler, SettableListenableFuture<WebSocketSession> connectFuture) {
        replacedert.notNull(request, "'request' is required");
        replacedert.notNull(handler, "'handler' is required");
        replacedert.notNull(connectFuture, "'connectFuture' is required");
        this.request = request;
        this.webSocketHandler = handler;
        this.connectFuture = connectFuture;
    }

    @Override
    public String getId() {
        return this.request.getSockJsUrlInfo().getSessionId();
    }

    @Override
    public URI getUri() {
        return this.request.getSockJsUrlInfo().getSockJsUrl();
    }

    @Override
    public HttpHeaders getHandshakeHeaders() {
        return this.request.getHandshakeHeaders();
    }

    @Override
    public Map<String, Object> getAttributes() {
        return this.attributes;
    }

    @Override
    public Principal getPrincipal() {
        return this.request.getUser();
    }

    public SockJsMessageCodec getMessageCodec() {
        return this.request.getMessageCodec();
    }

    public WebSocketHandler getWebSocketHandler() {
        return this.webSocketHandler;
    }

    /**
     * Return a timeout cleanup task to invoke if the SockJS sessions is not
     * fully established within the retransmission timeout period calculated in
     * {@code SockJsRequest} based on the duration of the initial SockJS "Info"
     * request.
     */
    Runnable getTimeoutTask() {
        return new Runnable() {

            @Override
            public void run() {
                closeInternal(new CloseStatus(2007, "Transport timed out"));
            }
        };
    }

    @Override
    public boolean isOpen() {
        return State.OPEN.equals(this.state);
    }

    public boolean isDisconnected() {
        return (State.CLOSING.equals(this.state) || State.CLOSED.equals(this.state));
    }

    @Override
    public final void sendMessage(WebSocketMessage<?> message) throws IOException {
        replacedert.state(State.OPEN.equals(this.state), this + " is not open, current state=" + this.state);
        replacedert.isInstanceOf(TextMessage.clreplaced, message, this + " supports text messages only.");
        String payload = ((TextMessage) message).getPayload();
        payload = getMessageCodec().encode(new String[] { payload });
        // the client-side doesn't need message framing (letter "a")
        payload = payload.substring(1);
        message = new TextMessage(payload);
        if (logger.isTraceEnabled()) {
            logger.trace("Sending message " + message + " in " + this);
        }
        sendInternal((TextMessage) message);
    }

    protected abstract void sendInternal(TextMessage textMessage) throws IOException;

    @Override
    public final void close() throws IOException {
        close(CloseStatus.NORMAL);
    }

    @Override
    public final void close(CloseStatus status) {
        replacedert.isTrue(status != null && isUserSetStatus(status), "Invalid close status: " + status);
        if (logger.isDebugEnabled()) {
            logger.debug("Closing session with " + status + " in " + this);
        }
        closeInternal(status);
    }

    private boolean isUserSetStatus(CloseStatus status) {
        return (status.getCode() == 1000 || (status.getCode() >= 3000 && status.getCode() <= 4999));
    }

    protected void closeInternal(CloseStatus status) {
        if (this.state == null) {
            logger.warn("Ignoring close since connect() was never invoked");
            return;
        }
        if (State.CLOSING.equals(this.state) || State.CLOSED.equals(this.state)) {
            logger.debug("Ignoring close (already closing or closed), current state=" + this.state);
            return;
        }
        this.state = State.CLOSING;
        this.closeStatus = status;
        try {
            disconnect(status);
        } catch (Throwable ex) {
            if (logger.isErrorEnabled()) {
                logger.error("Failed to close " + this, ex);
            }
        }
    }

    protected abstract void disconnect(CloseStatus status) throws IOException;

    public void handleFrame(String payload) {
        SockJsFrame frame = new SockJsFrame(payload);
        if (SockJsFrameType.OPEN.equals(frame.getType())) {
            handleOpenFrame();
        } else if (SockJsFrameType.MESSAGE.equals(frame.getType())) {
            handleMessageFrame(frame);
        } else if (SockJsFrameType.CLOSE.equals(frame.getType())) {
            handleCloseFrame(frame);
        } else if (SockJsFrameType.HEARTBEAT.equals(frame.getType())) {
            if (logger.isTraceEnabled()) {
                logger.trace("Received heartbeat in " + this);
            }
        } else {
            // should never happen
            throw new IllegalStateException("Unknown SockJS frame type " + frame + " in " + this);
        }
    }

    private void handleOpenFrame() {
        if (logger.isDebugEnabled()) {
            logger.debug("Processing SockJS open frame in " + this);
        }
        if (State.NEW.equals(state)) {
            this.state = State.OPEN;
            try {
                this.webSocketHandler.afterConnectionEstablished(this);
                this.connectFuture.set(this);
            } catch (Throwable ex) {
                if (logger.isErrorEnabled()) {
                    Clreplaced<?> type = this.webSocketHandler.getClreplaced();
                    logger.error(type + ".afterConnectionEstablished threw exception in " + this, ex);
                }
            }
        } else {
            if (logger.isDebugEnabled()) {
                logger.debug("Open frame received in " + getId() + " but we're not" + "connecting (current state=" + this.state + "). The server might " + "have been restarted and lost track of the session.");
            }
            closeInternal(new CloseStatus(1006, "Server lost session"));
        }
    }

    private void handleMessageFrame(SockJsFrame frame) {
        if (!isOpen()) {
            if (logger.isErrorEnabled()) {
                logger.error("Ignoring received message due to state=" + this.state + " in " + this);
            }
            return;
        }
        String[] messages;
        try {
            messages = getMessageCodec().decode(frame.getFrameData());
        } catch (IOException ex) {
            if (logger.isErrorEnabled()) {
                logger.error("Failed to decode data for SockJS \"message\" frame: " + frame + " in " + this, ex);
            }
            closeInternal(CloseStatus.BAD_DATA);
            return;
        }
        if (logger.isTraceEnabled()) {
            logger.trace("Processing SockJS message frame " + frame.getContent() + " in " + this);
        }
        for (String message : messages) {
            try {
                if (isOpen()) {
                    this.webSocketHandler.handleMessage(this, new TextMessage(message));
                }
            } catch (Throwable ex) {
                Clreplaced<?> type = this.webSocketHandler.getClreplaced();
                logger.error(type + ".handleMessage threw an exception on " + frame + " in " + this, ex);
            }
        }
    }

    private void handleCloseFrame(SockJsFrame frame) {
        CloseStatus closeStatus = CloseStatus.NO_STATUS_CODE;
        try {
            String[] data = getMessageCodec().decode(frame.getFrameData());
            if (data.length == 2) {
                closeStatus = new CloseStatus(Integer.valueOf(data[0]), data[1]);
            }
            if (logger.isDebugEnabled()) {
                logger.debug("Processing SockJS close frame with " + closeStatus + " in " + this);
            }
        } catch (IOException ex) {
            if (logger.isErrorEnabled()) {
                logger.error("Failed to decode data for " + frame + " in " + this, ex);
            }
        }
        closeInternal(closeStatus);
    }

    public void handleTransportError(Throwable error) {
        try {
            if (logger.isErrorEnabled()) {
                logger.error("Transport error in " + this, error);
            }
            this.webSocketHandler.handleTransportError(this, error);
        } catch (Exception ex) {
            Clreplaced<?> type = this.webSocketHandler.getClreplaced();
            if (logger.isErrorEnabled()) {
                logger.error(type + ".handleTransportError threw an exception", ex);
            }
        }
    }

    public void afterTransportClosed(CloseStatus closeStatus) {
        this.closeStatus = (this.closeStatus != null ? this.closeStatus : closeStatus);
        replacedert.state(this.closeStatus != null, "CloseStatus not available");
        if (logger.isDebugEnabled()) {
            logger.debug("Transport closed with " + this.closeStatus + " in " + this);
        }
        this.state = State.CLOSED;
        try {
            this.webSocketHandler.afterConnectionClosed(this, this.closeStatus);
        } catch (Exception ex) {
            if (logger.isErrorEnabled()) {
                Clreplaced<?> type = this.webSocketHandler.getClreplaced();
                logger.error(type + ".afterConnectionClosed threw an exception", ex);
            }
        }
    }

    @Override
    public String toString() {
        return getClreplaced().getSimpleName() + "[id='" + getId() + ", url=" + getUri() + "]";
    }

    private enum State {

        NEW, OPEN, CLOSING, CLOSED
    }
}

@Override
public ListenableFuture<ClientHttpResponse> executeAsync() throws IOException {
    SettableListenableFuture<ClientHttpResponse> future = new SettableListenableFuture<ClientHttpResponse>();
    future.set(execute());
    return future;
}

@Before
public void setup() {
    bigQueryTemplate = mock(BigQueryTemplate.clreplaced);
    SettableListenableFuture<Job> result = new SettableListenableFuture<>();
    result.set(mock(Job.clreplaced));
    when(bigQueryTemplate.writeDataToTable(any(), any(), any(), any())).thenReturn(result);
    messageHandler = new BigQueryFileMessageHandler(bigQueryTemplate);
}

@Override
public final ListenableFuture<WebSocketSession> doHandshake(WebSocketHandler handler, @Nullable WebSocketHttpHeaders headers, URI url) {
    replacedert.notNull(handler, "WebSocketHandler is required");
    replacedert.notNull(url, "URL is required");
    String scheme = url.getScheme();
    if (!supportedProtocols.contains(scheme)) {
        throw new IllegalArgumentException("Invalid scheme: '" + scheme + "'");
    }
    SettableListenableFuture<WebSocketSession> connectFuture = new SettableListenableFuture<>();
    try {
        SockJsUrlInfo sockJsUrlInfo = new SockJsUrlInfo(url);
        ServerInfo serverInfo = getServerInfo(sockJsUrlInfo, getHttpRequestHeaders(headers));
        createRequest(sockJsUrlInfo, headers, serverInfo).connect(handler, connectFuture);
    } catch (Throwable exception) {
        if (logger.isErrorEnabled()) {
            logger.error("Initial SockJS \"Info\" request to server failed, url=" + url, exception);
        }
        connectFuture.setException(exception);
    }
    return connectFuture;
}

// Transport methods
@Override
public ListenableFuture<WebSocketSession> connect(TransportRequest request, WebSocketHandler handler) {
    SettableListenableFuture<WebSocketSession> connectFuture = new SettableListenableFuture<>();
    XhrClientSockJsSession session = new XhrClientSockJsSession(request, handler, this, connectFuture);
    request.addTimeoutTask(session.getTimeoutTask());
    URI receiveUrl = request.getTransportUrl();
    if (logger.isDebugEnabled()) {
        logger.debug("Starting XHR " + (isXhrStreamingDisabled() ? "Polling" : "Streaming") + "session url=" + receiveUrl);
    }
    HttpHeaders handshakeHeaders = new HttpHeaders();
    handshakeHeaders.putAll(request.getHandshakeHeaders());
    connectInternal(request, handler, receiveUrl, handshakeHeaders, session, connectFuture);
    return connectFuture;
}

private SettableListenableFuture<Job> createJobFuture(Job pendingJob) {
    // Prepare the polling task for the ListenableFuture result returned to end-user
    SettableListenableFuture<Job> result = new SettableListenableFuture<>();
    ScheduledFuture<?> scheduledFuture = taskScheduler.scheduleAtFixedRate(() -> {
        try {
            Job job = pendingJob.reload();
            if (State.DONE.equals(job.getStatus().getState())) {
                if (job.getStatus().getError() != null) {
                    result.setException(new BigQueryException(job.getStatus().getError().getMessage()));
                } else {
                    result.set(job);
                }
            }
        } catch (Exception e) {
            result.setException(new BigQueryException(e.getMessage()));
        }
    }, this.jobPollInterval);
    result.addCallback(response -> scheduledFuture.cancel(true), response -> {
        pendingJob.cancel();
        scheduledFuture.cancel(true);
    });
    return result;
}

@Override
public final ListenableFuture<WebSocketSession> doHandshake(WebSocketHandler handler, @Nullable WebSocketHttpHeaders headers, URI url) {
    replacedert.notNull(handler, "WebSocketHandler is required");
    replacedert.notNull(url, "URL is required");
    String scheme = url.getScheme();
    if (!supportedProtocols.contains(scheme)) {
        throw new IllegalArgumentException("Invalid scheme: '" + scheme + "'");
    }
    SettableListenableFuture<WebSocketSession> connectFuture = new SettableListenableFuture<>();
    try {
        SockJsUrlInfo sockJsUrlInfo = new SockJsUrlInfo(url);
        ServerInfo serverInfo = getServerInfo(sockJsUrlInfo, getHttpRequestHeaders(headers));
        createRequest(sockJsUrlInfo, headers, serverInfo).connect(handler, connectFuture);
    } catch (Exception exception) {
        if (logger.isErrorEnabled()) {
            logger.error("Initial SockJS \"Info\" request to server failed, url=" + url, exception);
        }
        connectFuture.setException(exception);
    }
    return connectFuture;
}

public static <T> ListenableFuture<T> toListenable(CompletionStage<T> stage) {
    SettableListenableFuture<T> future = new SettableListenableFuture<>();
    stage.whenComplete((v, t) -> {
        if (t == null) {
            future.set(v);
        } else {
            future.setException(t);
        }
    });
    return future;
}

@Override
public final ListenableFuture<WebSocketSession> doHandshake(WebSocketHandler handler, WebSocketHttpHeaders headers, URI url) {
    replacedert.notNull(handler, "WebSocketHandler is required");
    replacedert.notNull(url, "URL is required");
    String scheme = url.getScheme();
    if (!supportedProtocols.contains(scheme)) {
        throw new IllegalArgumentException("Invalid scheme: '" + scheme + "'");
    }
    SettableListenableFuture<WebSocketSession> connectFuture = new SettableListenableFuture<WebSocketSession>();
    try {
        SockJsUrlInfo sockJsUrlInfo = new SockJsUrlInfo(url);
        ServerInfo serverInfo = getServerInfo(sockJsUrlInfo, getHttpRequestHeaders(headers));
        createRequest(sockJsUrlInfo, headers, serverInfo).connect(handler, connectFuture);
    } catch (Throwable exception) {
        if (logger.isErrorEnabled()) {
            logger.error("Initial SockJS \"Info\" request to server failed, url=" + url, exception);
        }
        connectFuture.setException(exception);
    }
    return connectFuture;
}

// Transport methods
@Override
@SuppressWarnings("deprecation")
public ListenableFuture<WebSocketSession> connect(TransportRequest request, WebSocketHandler handler) {
    SettableListenableFuture<WebSocketSession> connectFuture = new SettableListenableFuture<WebSocketSession>();
    XhrClientSockJsSession session = new XhrClientSockJsSession(request, handler, this, connectFuture);
    request.addTimeoutTask(session.getTimeoutTask());
    URI receiveUrl = request.getTransportUrl();
    if (logger.isDebugEnabled()) {
        logger.debug("Starting XHR " + (isXhrStreamingDisabled() ? "Polling" : "Streaming") + "session url=" + receiveUrl);
    }
    HttpHeaders handshakeHeaders = new HttpHeaders();
    handshakeHeaders.putAll(getRequestHeaders());
    handshakeHeaders.putAll(request.getHandshakeHeaders());
    connectInternal(request, handler, receiveUrl, handshakeHeaders, session, connectFuture);
    return connectFuture;
}

@Override
protected ListenableFuture<ClientHttpResponse> executeInternal(final HttpHeaders headers) throws IOException {
    final SettableListenableFuture<ClientHttpResponse> responseFuture = new SettableListenableFuture<ClientHttpResponse>();
    ChannelFutureListener connectionListener = new ChannelFutureListener() {

        @Override
        public void operationComplete(ChannelFuture future) throws Exception {
            if (future.isSuccess()) {
                Channel channel = future.channel();
                channel.pipeline().addLast(new RequestExecuteHandler(responseFuture));
                FullHttpRequest nettyRequest = createFullHttpRequest(headers);
                channel.writeAndFlush(nettyRequest);
            } else {
                responseFuture.setException(future.cause());
            }
        }
    };
    this.bootstrap.connect(this.uri.getHost(), getPort(this.uri)).addListener(connectionListener);
    return responseFuture;
}

/**
 * Unit tests for {@link WebSocketStompClient}.
 *
 * @author Rossen Stoyanchev
 */
public clreplaced WebSocketStompClientTests {

    @Mock
    private TaskScheduler taskScheduler;

    @Mock
    private ConnectionHandlingStompSession stompSession;

    @Mock
    private WebSocketSession webSocketSession;

    private TestWebSocketStompClient stompClient;

    private ArgumentCaptor<WebSocketHandler> webSocketHandlerCaptor;

    private SettableListenableFuture<WebSocketSession> handshakeFuture;

    @Before
    public void setUp() throws Exception {
        MockitoAnnotations.initMocks(this);
        WebSocketClient webSocketClient = mock(WebSocketClient.clreplaced);
        this.stompClient = new TestWebSocketStompClient(webSocketClient);
        this.stompClient.setTaskScheduler(this.taskScheduler);
        this.stompClient.setStompSession(this.stompSession);
        this.webSocketHandlerCaptor = ArgumentCaptor.forClreplaced(WebSocketHandler.clreplaced);
        this.handshakeFuture = new SettableListenableFuture<>();
        given(webSocketClient.doHandshake(this.webSocketHandlerCaptor.capture(), any(), any(URI.clreplaced))).willReturn(this.handshakeFuture);
    }

    @Test
    public void webSocketHandshakeFailure() throws Exception {
        connect();
        IllegalStateException handshakeFailure = new IllegalStateException("simulated exception");
        this.handshakeFuture.setException(handshakeFailure);
        verify(this.stompSession).afterConnectFailure(same(handshakeFailure));
    }

    @Test
    public void webSocketConnectionEstablished() throws Exception {
        connect().afterConnectionEstablished(this.webSocketSession);
        verify(this.stompSession).afterConnected(notNull());
    }

    @Test
    public void webSocketTransportError() throws Exception {
        IllegalStateException exception = new IllegalStateException("simulated exception");
        connect().handleTransportError(this.webSocketSession, exception);
        verify(this.stompSession).handleFailure(same(exception));
    }

    @Test
    public void webSocketConnectionClosed() throws Exception {
        connect().afterConnectionClosed(this.webSocketSession, CloseStatus.NORMAL);
        verify(this.stompSession).afterConnectionClosed();
    }

    @Test
    @SuppressWarnings({ "unchecked", "rawtypes" })
    public void handleWebSocketMessage() throws Exception {
        String text = "SEND\na:alpha\n\nMessage payload\0";
        connect().handleMessage(this.webSocketSession, new TextMessage(text));
        ArgumentCaptor<Message> captor = ArgumentCaptor.forClreplaced(Message.clreplaced);
        verify(this.stompSession).handleMessage(captor.capture());
        Message<byte[]> message = captor.getValue();
        replacedertNotNull(message);
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.clreplaced);
        StompHeaders headers = StompHeaders.readOnlyStompHeaders(accessor.toNativeHeaderMap());
        replacedertEquals(StompCommand.SEND, accessor.getCommand());
        replacedertEquals("alpha", headers.getFirst("a"));
        replacedertEquals("Message payload", new String(message.getPayload(), StandardCharsets.UTF_8));
    }

    @Test
    @SuppressWarnings({ "unchecked", "rawtypes" })
    public void handleWebSocketMessageSplitAcrossTwoMessage() throws Exception {
        WebSocketHandler webSocketHandler = connect();
        String part1 = "SEND\na:alpha\n\nMessage";
        webSocketHandler.handleMessage(this.webSocketSession, new TextMessage(part1));
        verifyNoMoreInteractions(this.stompSession);
        String part2 = " payload\0";
        webSocketHandler.handleMessage(this.webSocketSession, new TextMessage(part2));
        ArgumentCaptor<Message> captor = ArgumentCaptor.forClreplaced(Message.clreplaced);
        verify(this.stompSession).handleMessage(captor.capture());
        Message<byte[]> message = captor.getValue();
        replacedertNotNull(message);
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.clreplaced);
        StompHeaders headers = StompHeaders.readOnlyStompHeaders(accessor.toNativeHeaderMap());
        replacedertEquals(StompCommand.SEND, accessor.getCommand());
        replacedertEquals("alpha", headers.getFirst("a"));
        replacedertEquals("Message payload", new String(message.getPayload(), StandardCharsets.UTF_8));
    }

    @Test
    @SuppressWarnings({ "unchecked", "rawtypes" })
    public void handleWebSocketMessageBinary() throws Exception {
        String text = "SEND\na:alpha\n\nMessage payload\0";
        connect().handleMessage(this.webSocketSession, new BinaryMessage(text.getBytes(StandardCharsets.UTF_8)));
        ArgumentCaptor<Message> captor = ArgumentCaptor.forClreplaced(Message.clreplaced);
        verify(this.stompSession).handleMessage(captor.capture());
        Message<byte[]> message = captor.getValue();
        replacedertNotNull(message);
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.clreplaced);
        StompHeaders headers = StompHeaders.readOnlyStompHeaders(accessor.toNativeHeaderMap());
        replacedertEquals(StompCommand.SEND, accessor.getCommand());
        replacedertEquals("alpha", headers.getFirst("a"));
        replacedertEquals("Message payload", new String(message.getPayload(), StandardCharsets.UTF_8));
    }

    @Test
    public void handleWebSocketMessagePong() throws Exception {
        connect().handleMessage(this.webSocketSession, new PongMessage());
        verifyNoMoreInteractions(this.stompSession);
    }

    @Test
    public void sendWebSocketMessage() throws Exception {
        StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.SEND);
        accessor.setDestination("/topic/foo");
        byte[] payload = "payload".getBytes(StandardCharsets.UTF_8);
        getTcpConnection().send(MessageBuilder.createMessage(payload, accessor.getMessageHeaders()));
        ArgumentCaptor<TextMessage> textMessageCaptor = ArgumentCaptor.forClreplaced(TextMessage.clreplaced);
        verify(this.webSocketSession).sendMessage(textMessageCaptor.capture());
        TextMessage textMessage = textMessageCaptor.getValue();
        replacedertNotNull(textMessage);
        replacedertEquals("SEND\ndestination:/topic/foo\ncontent-length:7\n\npayload\0", textMessage.getPayload());
    }

    @Test
    public void sendWebSocketBinary() throws Exception {
        StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.SEND);
        accessor.setDestination("/b");
        accessor.setContentType(MimeTypeUtils.APPLICATION_OCTET_STREAM);
        byte[] payload = "payload".getBytes(StandardCharsets.UTF_8);
        getTcpConnection().send(MessageBuilder.createMessage(payload, accessor.getMessageHeaders()));
        ArgumentCaptor<BinaryMessage> binaryMessageCaptor = ArgumentCaptor.forClreplaced(BinaryMessage.clreplaced);
        verify(this.webSocketSession).sendMessage(binaryMessageCaptor.capture());
        BinaryMessage binaryMessage = binaryMessageCaptor.getValue();
        replacedertNotNull(binaryMessage);
        replacedertEquals("SEND\ndestination:/b\ncontent-type:application/octet-stream\ncontent-length:7\n\npayload\0", new String(binaryMessage.getPayload().array(), StandardCharsets.UTF_8));
    }

    @Test
    public void heartbeatDefaultValue() throws Exception {
        WebSocketStompClient stompClient = new WebSocketStompClient(mock(WebSocketClient.clreplaced));
        replacedertArrayEquals(new long[] { 0, 0 }, stompClient.getDefaultHeartbeat());
        StompHeaders connectHeaders = stompClient.processConnectHeaders(null);
        replacedertArrayEquals(new long[] { 0, 0 }, connectHeaders.getHeartbeat());
    }

    @Test
    public void heartbeatDefaultValueWithScheduler() throws Exception {
        WebSocketStompClient stompClient = new WebSocketStompClient(mock(WebSocketClient.clreplaced));
        stompClient.setTaskScheduler(mock(TaskScheduler.clreplaced));
        replacedertArrayEquals(new long[] { 10000, 10000 }, stompClient.getDefaultHeartbeat());
        StompHeaders connectHeaders = stompClient.processConnectHeaders(null);
        replacedertArrayEquals(new long[] { 10000, 10000 }, connectHeaders.getHeartbeat());
    }

    @Test
    public void heartbeatDefaultValueSetWithoutScheduler() throws Exception {
        WebSocketStompClient stompClient = new WebSocketStompClient(mock(WebSocketClient.clreplaced));
        stompClient.setDefaultHeartbeat(new long[] { 5, 5 });
        try {
            stompClient.processConnectHeaders(null);
            fail("Expected IllegalStateException");
        } catch (IllegalStateException ex) {
        // ignore
        }
    }

    @Test
    public void readInactivityAfterDelayHasElapsed() throws Exception {
        TcpConnection<byte[]> tcpConnection = getTcpConnection();
        Runnable runnable = mock(Runnable.clreplaced);
        long delay = 2;
        tcpConnection.onReadInactivity(runnable, delay);
        testInactivityTaskScheduling(runnable, delay, 10);
    }

    @Test
    public void readInactivityBeforeDelayHasElapsed() throws Exception {
        TcpConnection<byte[]> tcpConnection = getTcpConnection();
        Runnable runnable = mock(Runnable.clreplaced);
        long delay = 10000;
        tcpConnection.onReadInactivity(runnable, delay);
        testInactivityTaskScheduling(runnable, delay, 0);
    }

    @Test
    public void writeInactivityAfterDelayHasElapsed() throws Exception {
        TcpConnection<byte[]> tcpConnection = getTcpConnection();
        Runnable runnable = mock(Runnable.clreplaced);
        long delay = 2;
        tcpConnection.onWriteInactivity(runnable, delay);
        testInactivityTaskScheduling(runnable, delay, 10);
    }

    @Test
    public void writeInactivityBeforeDelayHasElapsed() throws Exception {
        TcpConnection<byte[]> tcpConnection = getTcpConnection();
        Runnable runnable = mock(Runnable.clreplaced);
        long delay = 1000;
        tcpConnection.onWriteInactivity(runnable, delay);
        testInactivityTaskScheduling(runnable, delay, 0);
    }

    @Test
    @SuppressWarnings({ "rawtypes", "unchecked" })
    public void cancelInactivityTasks() throws Exception {
        TcpConnection<byte[]> tcpConnection = getTcpConnection();
        ScheduledFuture future = mock(ScheduledFuture.clreplaced);
        given(this.taskScheduler.scheduleWithFixedDelay(any(), eq(1L))).willReturn(future);
        tcpConnection.onReadInactivity(mock(Runnable.clreplaced), 2L);
        tcpConnection.onWriteInactivity(mock(Runnable.clreplaced), 2L);
        this.webSocketHandlerCaptor.getValue().afterConnectionClosed(this.webSocketSession, CloseStatus.NORMAL);
        verify(future, times(2)).cancel(true);
        verifyNoMoreInteractions(future);
    }

    private WebSocketHandler connect() {
        this.stompClient.connect("/foo", mock(StompSessionHandler.clreplaced));
        verify(this.stompSession).getSessionFuture();
        verifyNoMoreInteractions(this.stompSession);
        WebSocketHandler webSocketHandler = this.webSocketHandlerCaptor.getValue();
        replacedertNotNull(webSocketHandler);
        return webSocketHandler;
    }

    @SuppressWarnings("unchecked")
    private TcpConnection<byte[]> getTcpConnection() throws Exception {
        WebSocketHandler webSocketHandler = connect();
        webSocketHandler.afterConnectionEstablished(this.webSocketSession);
        return (TcpConnection<byte[]>) webSocketHandler;
    }

    private void testInactivityTaskScheduling(Runnable runnable, long delay, long sleepTime) throws InterruptedException {
        ArgumentCaptor<Runnable> inactivityTaskCaptor = ArgumentCaptor.forClreplaced(Runnable.clreplaced);
        verify(this.taskScheduler).scheduleWithFixedDelay(inactivityTaskCaptor.capture(), eq(delay / 2));
        verifyNoMoreInteractions(this.taskScheduler);
        if (sleepTime > 0) {
            Thread.sleep(sleepTime);
        }
        Runnable inactivityTask = inactivityTaskCaptor.getValue();
        replacedertNotNull(inactivityTask);
        inactivityTask.run();
        if (sleepTime > 0) {
            verify(runnable).run();
        } else {
            verifyNoMoreInteractions(runnable);
        }
    }

    private static clreplaced TestWebSocketStompClient extends WebSocketStompClient {

        private ConnectionHandlingStompSession stompSession;

        public TestWebSocketStompClient(WebSocketClient webSocketClient) {
            super(webSocketClient);
        }

        public void setStompSession(ConnectionHandlingStompSession stompSession) {
            this.stompSession = stompSession;
        }

        @Override
        protected ConnectionHandlingStompSession createSession(StompHeaders headers, StompSessionHandler handler) {
            return this.stompSession;
        }
    }
}

See More Examples