Examples with DrainReader com.linkedin.r2.message.stream.entitystream.DrainReader used on opensource projects

Example 6 with DrainReader

use of com.linkedin.r2.message.stream.entitystream.DrainReader in project rest.li by linkedin.

the class TestClientTimeout method testReadAfterTimeout.

@Test(dataProvider = "clients")
public void testReadAfterTimeout(Client client) throws Exception {
    StreamRequest request = new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, NORMAL_URI)).build(EntityStreams.emptyStream());
    final CountDownLatch latch = new CountDownLatch(1);
    final AtomicReference<StreamResponse> response = new AtomicReference<StreamResponse>();
    client.streamRequest(request, new Callback<StreamResponse>() {

        @Override
        public void onError(Throwable e) {
            latch.countDown();
        }

        @Override
        public void onSuccess(StreamResponse result) {
            response.set(result);
            latch.countDown();
        }
    });
    latch.await(5000, TimeUnit.MILLISECONDS);
    Assert.assertNotNull(response.get());
    // let it timeout before we read
    Thread.sleep(5000);
    final AtomicReference<Throwable> throwable = new AtomicReference<Throwable>();
    final CountDownLatch errorLatch = new CountDownLatch(1);
    Reader reader = new DrainReader() {

        @Override
        public void onError(Throwable ex) {
            throwable.set(ex);
            errorLatch.countDown();
        }
    };
    response.get().getEntityStream().setReader(reader);
    errorLatch.await(5000, TimeUnit.MILLISECONDS);
    Assert.assertNotNull(throwable.get());
    Throwable rootCause = ExceptionUtils.getRootCause(throwable.get());
    Assert.assertTrue(rootCause instanceof TimeoutException);
    Assert.assertEquals(rootCause.getMessage(), "Timeout while receiving the response entity.");
}

Example 7 with DrainReader

use of com.linkedin.r2.message.stream.entitystream.DrainReader in project rest.li by linkedin.

the class TestJetty404 method testJetty404.

// make sure jetty's default behavior will read all the request bytes in case of 404
@Test
public void testJetty404() throws Exception {
    BytesWriter writer = new BytesWriter(200 * 1024, (byte) 100);
    final AtomicReference<Throwable> exRef = new AtomicReference<Throwable>();
    final CountDownLatch latch = new CountDownLatch(1);
    _client.streamRequest(new StreamRequestBuilder(Bootstrap.createHttpURI(PORT, URI.create("/wrong-path"))).build(EntityStreams.newEntityStream(writer)), new Callback<StreamResponse>() {

        @Override
        public void onError(Throwable e) {
            exRef.set(e);
            latch.countDown();
        }

        @Override
        public void onSuccess(StreamResponse result) {
            latch.countDown();
        }
    });
    latch.await(5000, TimeUnit.MILLISECONDS);
    Assert.assertTrue(writer.isDone());
    Throwable ex = exRef.get();
    Assert.assertTrue(ex instanceof StreamException);
    StreamResponse response = ((StreamException) ex).getResponse();
    Assert.assertEquals(response.getStatus(), RestStatus.NOT_FOUND);
    response.getEntityStream().setReader(new DrainReader());
}

Example 8 with DrainReader

use of com.linkedin.r2.message.stream.entitystream.DrainReader in project rest.li by linkedin.

the class TestJetty404 method setup.

@BeforeClass
public void setup() throws IOException {
    _clientFactory = new HttpClientFactory();
    _client = new TransportClientAdapter(_clientFactory.getClient(Collections.<String, String>emptyMap()), true);
    _server = new HttpServerFactory().createH2cServer(PORT, "/correct-path", 50, new TransportDispatcher() {

        @Override
        public void handleRestRequest(RestRequest req, Map<String, String> wireAttrs, RequestContext requestContext, TransportCallback<RestResponse> callback) {
            callback.onResponse(TransportResponseImpl.success(new RestResponseBuilder().build()));
        }

        @Override
        public void handleStreamRequest(StreamRequest req, Map<String, String> wireAttrs, RequestContext requestContext, TransportCallback<StreamResponse> callback) {
            req.getEntityStream().setReader(new DrainReader());
            callback.onResponse(TransportResponseImpl.success(new StreamResponseBuilder().build(EntityStreams.emptyStream())));
        }
    }, true);
    _server.start();
}