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

Example 11 with WriteHandle

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

the class TestMIMEInputStream method testMultipleOnWritePossibleDataSources.

@Test(dataProvider = "multipleOnWritePossibleDataSources")
public void testMultipleOnWritePossibleDataSources(final byte[] inputData, final StrictByteArrayInputStream inputStream, final int onWritePossibles, final int writesRemainingPerOnWritePossible, final int expectedTotalWrites, final int expectedWriteHandleRemainingCalls) {
    // Setup:
    final WriteHandle writeHandle = Mockito.mock(WriteHandle.class);
    final MultiPartMIMEInputStream multiPartMIMEInputStream = new MultiPartMIMEInputStream.Builder(inputStream, _scheduledExecutorService, Collections.<String, String>emptyMap()).withWriteChunkSize(TEST_CHUNK_SIZE).build();
    // We want to simulate a decreasing return from writeHandle.remaining().
    // Note that the 0 is added on later so we stop at 1:
    final Integer[] remainingWriteHandleCount;
    if (writesRemainingPerOnWritePossible > 1) {
        // This represents writeHandle.remaining() -> n, n -1, n - 2, .... 1, 0
        remainingWriteHandleCount = new Integer[writesRemainingPerOnWritePossible - 1];
        int writeHandleCountTemp = writesRemainingPerOnWritePossible;
        for (int i = 0; i < writesRemainingPerOnWritePossible - 1; i++) {
            remainingWriteHandleCount[i] = --writeHandleCountTemp;
        }
    } else {
        // This represents writeHandle.remaining() -> 1, 0
        remainingWriteHandleCount = new Integer[] {};
    }
    OngoingStubbing<Integer> writeHandleOngoingStubbing = when(writeHandle.remaining());
    for (int i = 0; i < onWritePossibles; i++) {
        // Each onWritePossible() will provide the corresponding number of writes remaining on the write handle.
        // When the writeHandle.remaining() reaches zero, onWritePossible() will be invoked again.
        // Mockito does not mention that chaining requires keeping the references and only allows one append at a time.
        // A painful lesson that I had to learn so I will leave a comment here for future people.
        writeHandleOngoingStubbing = writeHandleOngoingStubbing.thenReturn(writesRemainingPerOnWritePossible, remainingWriteHandleCount);
        writeHandleOngoingStubbing = writeHandleOngoingStubbing.thenAnswer(new Answer<Integer>() {

            @Override
            public Integer answer(InvocationOnMock invocation) throws Throwable {
                // There is no better way to do this with mockito. R2 observes that 0 has been returned and THEN
                // invokes onWritePossible(). We need to make sure that the value of 0 is returned
                // and then onWritePossible() is invoked afterwards.
                _scheduledExecutorService.schedule(new Runnable() {

                    @Override
                    public void run() {
                        multiPartMIMEInputStream.onWritePossible();
                    }
                }, 1000, TimeUnit.MILLISECONDS);
                return 0;
            }
        });
    }
    final ByteArrayOutputStream byteArrayOutputStream = setupMockWriteHandleToOutputStream(writeHandle);
    // Setup for done()
    final CountDownLatch doneLatch = new CountDownLatch(1);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            doneLatch.countDown();
            return null;
        }
    }).when(writeHandle).done();
    // /////////////////////////////////
    // Start things off
    // Init the data source
    multiPartMIMEInputStream.onInit(writeHandle);
    multiPartMIMEInputStream.onWritePossible();
    // Wait to finish
    try {
        boolean successful = doneLatch.await(_testTimeout, TimeUnit.MILLISECONDS);
        if (!successful) {
            Assert.fail("Timeout when waiting for input stream to completely transfer");
        }
    } catch (Exception exception) {
        Assert.fail("Unexpected exception when waiting for input stream to completely transfer");
    }
    // /////////////////////////////////
    // Assert
    Assert.assertEquals(byteArrayOutputStream.toByteArray(), inputData, "All data from the input stream should have successfully been transferred");
    Assert.assertEquals(inputStream.isClosed(), true);
    // Mock verifies:
    // The amount of times we write and the amount of times we call remaining() is the same.
    verify(writeHandle, times(expectedTotalWrites)).write(isA(ByteString.class));
    verify(writeHandle, times(expectedWriteHandleRemainingCalls)).remaining();
    verify(writeHandle, never()).error(isA(Throwable.class));
    verify(writeHandle, times(1)).done();
    verifyNoMoreInteractions(writeHandle);
}

Example 12 with WriteHandle

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

the class TestMIMEInputStream method abortWhenNoOutstandingReadTask.

// Abort in the middle of a write task.
@Test
public void abortWhenNoOutstandingReadTask() throws Exception {
    final StringBuilder builder = new StringBuilder();
    for (int i = 0; i < (TEST_CHUNK_SIZE * 10) + 2; i++) {
        builder.append('a');
    }
    // The slow byte array input stream will verify that we call an abort before the first read task is finished.
    final byte[] largeInputData = builder.toString().getBytes();
    final SlowByteArrayInputStream inputStream = new SlowByteArrayInputStream(largeInputData, 300, 10);
    final SlowByteArrayInputStream spyInputStream = spy(inputStream);
    // Setup:
    final WriteHandle writeHandle = Mockito.mock(WriteHandle.class);
    final MultiPartMIMEInputStream multiPartMIMEInputStream = new MultiPartMIMEInputStream.Builder(spyInputStream, _scheduledExecutorService, Collections.<String, String>emptyMap()).withWriteChunkSize(TEST_CHUNK_SIZE).build();
    // By the time the first onWritePossible() completes, half the data should be transferred
    // Then the abort task will run.
    when(writeHandle.remaining()).thenReturn(5, new Integer[] { 4, 3, 2, 1, 0 });
    final ByteArrayOutputStream byteArrayOutputStream = setupMockWriteHandleToOutputStream(writeHandle);
    // Setup for the close on the input stream.
    // The close must happen for the test to finish.
    final CountDownLatch closeInputStreamLatch = new CountDownLatch(1);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            closeInputStreamLatch.countDown();
            return null;
        }
    }).when(spyInputStream).close();
    // /////////////////////////////////
    // Start things off
    // Init the data source
    multiPartMIMEInputStream.onInit(writeHandle);
    multiPartMIMEInputStream.onWritePossible();
    multiPartMIMEInputStream.onAbort(new IOException());
    // Wait to finish
    try {
        boolean successful = closeInputStreamLatch.await(_testTimeout, TimeUnit.MILLISECONDS);
        if (!successful) {
            Assert.fail("Timeout when waiting for abort to happen!");
        }
    } catch (Exception exception) {
        Assert.fail("Unexpected exception when waiting for input stream to be closed!");
    }
    // /////////////////////////////////
    // Assert
    final byte[] expectedBytes = Arrays.copyOf(largeInputData, TEST_CHUNK_SIZE * 5);
    Assert.assertEquals(byteArrayOutputStream.toByteArray(), expectedBytes, "Partial data from the input stream should have successfully been transferred");
    // Mock verifies:
    verify(spyInputStream, times(1)).close();
    verify(spyInputStream, times(5)).read(isA(byte[].class));
    verify(writeHandle, times(5)).write(isA(ByteString.class));
    verify(writeHandle, times(6)).remaining();
    verify(writeHandle, never()).error(isA(Throwable.class));
    verify(writeHandle, never()).done();
    verifyNoMoreInteractions(writeHandle);
}

Example 13 with WriteHandle

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

the class TestMIMEInputStream method testTimeoutDataSources.

@Test(dataProvider = "timeoutDataSources")
public void testTimeoutDataSources(final TimeoutByteArrayInputStream timeoutByteArrayInputStream, final int expectedTotalWrites, final int expectedWriteHandleRemainingCalls, final byte[] expectedDataWritten, final CountDownLatch latch) {
    // Setup
    final WriteHandle writeHandle = Mockito.mock(WriteHandle.class);
    // For the maximum blocking time, we choose some value that isn't too short for a read to occur from the in memory
    // InputStream, but also not too long to prevent the test from taking a while to finish (due to waiting for the timeout
    // to occur).
    final MultiPartMIMEInputStream multiPartMIMEInputStream = new MultiPartMIMEInputStream.Builder(timeoutByteArrayInputStream, _scheduledExecutorService, Collections.<String, String>emptyMap()).withWriteChunkSize(TEST_CHUNK_SIZE).withMaximumBlockingTime(100).build();
    // Doesn't matter what we return here as long as its constant and above 0.
    when(writeHandle.remaining()).thenReturn(1);
    final ByteArrayOutputStream byteArrayOutputStream = setupMockWriteHandleToOutputStream(writeHandle);
    // Setup for error()
    final CountDownLatch errorLatch = new CountDownLatch(1);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            errorLatch.countDown();
            return null;
        }
    }).when(writeHandle).error(isA(Throwable.class));
    // /////////////////////////////////
    // Start things off
    // Init the data source
    multiPartMIMEInputStream.onInit(writeHandle);
    multiPartMIMEInputStream.onWritePossible();
    // Wait to finish
    try {
        boolean successful = errorLatch.await(_testTimeout, TimeUnit.MILLISECONDS);
        // Unblock the thread in the thread pool.
        latch.countDown();
        if (!successful) {
            Assert.fail("Timeout when waiting for input stream to completely transfer");
        }
    } catch (Exception exception) {
        Assert.fail("Unexpected exception when waiting for input stream to transfer");
    }
    // /////////////////////////////////
    // Assert
    Assert.assertEquals(byteArrayOutputStream.toByteArray(), expectedDataWritten, "Partial data should have been transferred in the case of a timeout");
    Assert.assertEquals(timeoutByteArrayInputStream.isClosed(), true);
    // Mock verifies:
    verify(writeHandle, times(expectedTotalWrites)).write(isA(ByteString.class));
    verify(writeHandle, times(expectedWriteHandleRemainingCalls)).remaining();
    // Since we can't override equals
    verify(writeHandle, times(1)).error(isA(TimeoutException.class));
    verify(writeHandle, never()).done();
    verifyNoMoreInteractions(writeHandle);
}