Examples with InvocationOnMock org.mockito.invocation.InvocationOnMock used on opensource projects

Example 11 with InvocationOnMock

use of org.mockito.invocation.InvocationOnMock in project flink by apache.

the class LocalInputChannelTest method testPartitionRequestExponentialBackoff.

@Test
public void testPartitionRequestExponentialBackoff() throws Exception {
    // Config
    Tuple2<Integer, Integer> backoff = new Tuple2<>(500, 3000);
    // Start with initial backoff, then keep doubling, and cap at max.
    int[] expectedDelays = { backoff._1(), 1000, 2000, backoff._2() };
    // Setup
    SingleInputGate inputGate = mock(SingleInputGate.class);
    BufferProvider bufferProvider = mock(BufferProvider.class);
    when(inputGate.getBufferProvider()).thenReturn(bufferProvider);
    ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
    LocalInputChannel ch = createLocalInputChannel(inputGate, partitionManager, backoff);
    when(partitionManager.createSubpartitionView(eq(ch.partitionId), eq(0), eq(bufferProvider), any(BufferAvailabilityListener.class))).thenThrow(new PartitionNotFoundException(ch.partitionId));
    Timer timer = mock(Timer.class);
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            ((TimerTask) invocation.getArguments()[0]).run();
            return null;
        }
    }).when(timer).schedule(any(TimerTask.class), anyLong());
    // Initial request
    ch.requestSubpartition(0);
    verify(partitionManager).createSubpartitionView(eq(ch.partitionId), eq(0), eq(bufferProvider), any(BufferAvailabilityListener.class));
    // Request subpartition and verify that the actual requests are delayed.
    for (long expected : expectedDelays) {
        ch.retriggerSubpartitionRequest(timer, 0);
        verify(timer).schedule(any(TimerTask.class), eq(expected));
    }
    // Exception after backoff is greater than the maximum backoff.
    try {
        ch.retriggerSubpartitionRequest(timer, 0);
        ch.getNextBuffer();
        fail("Did not throw expected exception.");
    } catch (Exception expected) {
    }
}

Example 12 with InvocationOnMock

use of org.mockito.invocation.InvocationOnMock in project flink by apache.

the class RecordWriterTest method testClearBuffersAfterExceptionInPartitionWriter.

@Test
public void testClearBuffersAfterExceptionInPartitionWriter() throws Exception {
    NetworkBufferPool buffers = null;
    BufferPool bufferPool = null;
    try {
        buffers = new NetworkBufferPool(1, 1024, MemoryType.HEAP);
        bufferPool = spy(buffers.createBufferPool(1, Integer.MAX_VALUE));
        ResultPartitionWriter partitionWriter = mock(ResultPartitionWriter.class);
        when(partitionWriter.getBufferProvider()).thenReturn(checkNotNull(bufferPool));
        when(partitionWriter.getNumberOfOutputChannels()).thenReturn(1);
        // Recycle buffer and throw Exception
        doAnswer(new Answer<Void>() {

            @Override
            public Void answer(InvocationOnMock invocation) throws Throwable {
                Buffer buffer = (Buffer) invocation.getArguments()[0];
                buffer.recycle();
                throw new RuntimeException("Expected test Exception");
            }
        }).when(partitionWriter).writeBuffer(any(Buffer.class), anyInt());
        RecordWriter<IntValue> recordWriter = new RecordWriter<>(partitionWriter);
        try {
            // manual flush here doesn't test this case (see next).
            for (; ; ) {
                recordWriter.emit(new IntValue(0));
            }
        } catch (Exception e) {
            // Verify that the buffer is not part of the record writer state after a failure
            // to flush it out. If the buffer is still part of the record writer state, this
            // will fail, because the buffer has already been recycled. NOTE: The mock
            // partition writer needs to recycle the buffer to correctly test this.
            recordWriter.clearBuffers();
        }
        // Verify expected methods have been called
        verify(partitionWriter, times(1)).writeBuffer(any(Buffer.class), anyInt());
        verify(bufferPool, times(1)).requestBufferBlocking();
        try {
            // Verify that manual flushing correctly clears the buffer.
            recordWriter.emit(new IntValue(0));
            recordWriter.flush();
            Assert.fail("Did not throw expected test Exception");
        } catch (Exception e) {
            recordWriter.clearBuffers();
        }
        // Verify expected methods have been called
        verify(partitionWriter, times(2)).writeBuffer(any(Buffer.class), anyInt());
        verify(bufferPool, times(2)).requestBufferBlocking();
        try {
            // Verify that broadcast emit correctly clears the buffer.
            for (; ; ) {
                recordWriter.broadcastEmit(new IntValue(0));
            }
        } catch (Exception e) {
            recordWriter.clearBuffers();
        }
        // Verify expected methods have been called
        verify(partitionWriter, times(3)).writeBuffer(any(Buffer.class), anyInt());
        verify(bufferPool, times(3)).requestBufferBlocking();
        try {
            // Verify that end of super step correctly clears the buffer.
            recordWriter.emit(new IntValue(0));
            recordWriter.broadcastEvent(EndOfSuperstepEvent.INSTANCE);
            Assert.fail("Did not throw expected test Exception");
        } catch (Exception e) {
            recordWriter.clearBuffers();
        }
        // Verify expected methods have been called
        verify(partitionWriter, times(4)).writeBuffer(any(Buffer.class), anyInt());
        verify(bufferPool, times(4)).requestBufferBlocking();
        try {
            // Verify that broadcasting and event correctly clears the buffer.
            recordWriter.emit(new IntValue(0));
            recordWriter.broadcastEvent(new TestTaskEvent());
            Assert.fail("Did not throw expected test Exception");
        } catch (Exception e) {
            recordWriter.clearBuffers();
        }
        // Verify expected methods have been called
        verify(partitionWriter, times(5)).writeBuffer(any(Buffer.class), anyInt());
        verify(bufferPool, times(5)).requestBufferBlocking();
    } finally {
        if (bufferPool != null) {
            assertEquals(1, bufferPool.getNumberOfAvailableMemorySegments());
            bufferPool.lazyDestroy();
        }
        if (buffers != null) {
            assertEquals(1, buffers.getNumberOfAvailableMemorySegments());
            buffers.destroy();
        }
    }
}

Example 13 with InvocationOnMock

use of org.mockito.invocation.InvocationOnMock in project flink by apache.

the class CancelPartitionRequestTest method testDuplicateCancel.

@Test
public void testDuplicateCancel() throws Exception {
    NettyServerAndClient serverAndClient = null;
    try {
        final TestPooledBufferProvider outboundBuffers = new TestPooledBufferProvider(16);
        ResultPartitionManager partitions = mock(ResultPartitionManager.class);
        ResultPartitionID pid = new ResultPartitionID();
        final CountDownLatch sync = new CountDownLatch(1);
        final ResultSubpartitionView view = spy(new InfiniteSubpartitionView(outboundBuffers, sync));
        // Return infinite subpartition
        when(partitions.createSubpartitionView(eq(pid), eq(0), any(BufferProvider.class), any(BufferAvailabilityListener.class))).thenAnswer(new Answer<ResultSubpartitionView>() {

            @Override
            public ResultSubpartitionView answer(InvocationOnMock invocationOnMock) throws Throwable {
                BufferAvailabilityListener listener = (BufferAvailabilityListener) invocationOnMock.getArguments()[3];
                listener.notifyBuffersAvailable(Long.MAX_VALUE);
                return view;
            }
        });
        PartitionRequestProtocol protocol = new PartitionRequestProtocol(partitions, mock(TaskEventDispatcher.class), mock(NetworkBufferPool.class));
        serverAndClient = initServerAndClient(protocol);
        Channel ch = connect(serverAndClient);
        // Request for non-existing input channel => results in cancel request
        InputChannelID inputChannelId = new InputChannelID();
        ch.writeAndFlush(new PartitionRequest(pid, 0, inputChannelId)).await();
        // Wait for the notification
        if (!sync.await(TestingUtils.TESTING_DURATION().toMillis(), TimeUnit.MILLISECONDS)) {
            fail("Timed out after waiting for " + TestingUtils.TESTING_DURATION().toMillis() + " ms to be notified about cancelled partition.");
        }
        ch.writeAndFlush(new CancelPartitionRequest(inputChannelId)).await();
        ch.close();
        NettyTestUtil.awaitClose(ch);
        verify(view, times(1)).releaseAllResources();
        verify(view, times(0)).notifySubpartitionConsumed();
    } finally {
        shutdown(serverAndClient);
    }
}

Example 14 with InvocationOnMock

use of org.mockito.invocation.InvocationOnMock in project flink by apache.

the class ClientTransportErrorHandlingTest method testExceptionOnWrite.

/**
	 * Verifies that failed client requests via {@link PartitionRequestClient} are correctly
	 * attributed to the respective {@link RemoteInputChannel}.
	 */
@Test
public void testExceptionOnWrite() throws Exception {
    NettyProtocol protocol = new NettyProtocol() {

        @Override
        public ChannelHandler[] getServerChannelHandlers() {
            return new ChannelHandler[0];
        }

        @Override
        public ChannelHandler[] getClientChannelHandlers() {
            return new PartitionRequestProtocol(mock(ResultPartitionProvider.class), mock(TaskEventDispatcher.class), mock(NetworkBufferPool.class)).getClientChannelHandlers();
        }
    };
    // We need a real server and client in this test, because Netty's EmbeddedChannel is
    // not failing the ChannelPromise of failed writes.
    NettyServerAndClient serverAndClient = initServerAndClient(protocol, createConfig());
    Channel ch = connect(serverAndClient);
    PartitionRequestClientHandler handler = getClientHandler(ch);
    // Last outbound handler throws Exception after 1st write
    ch.pipeline().addFirst(new ChannelOutboundHandlerAdapter() {

        int writeNum = 0;

        @Override
        public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
            if (writeNum >= 1) {
                throw new RuntimeException("Expected test exception.");
            }
            writeNum++;
            ctx.write(msg, promise);
        }
    });
    PartitionRequestClient requestClient = new PartitionRequestClient(ch, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));
    // Create input channels
    RemoteInputChannel[] rich = new RemoteInputChannel[] { createRemoteInputChannel(), createRemoteInputChannel() };
    final CountDownLatch sync = new CountDownLatch(1);
    // Do this with explicit synchronization. Otherwise this is not robust against slow timings
    // of the callback (e.g. we cannot just verify that it was called once, because there is
    // a chance that we do this too early).
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            sync.countDown();
            return null;
        }
    }).when(rich[1]).onError(isA(LocalTransportException.class));
    // First request is successful
    ChannelFuture f = requestClient.requestSubpartition(new ResultPartitionID(), 0, rich[0], 0);
    assertTrue(f.await().isSuccess());
    // Second request is *not* successful
    f = requestClient.requestSubpartition(new ResultPartitionID(), 0, rich[1], 0);
    assertFalse(f.await().isSuccess());
    // Only the second channel should be notified about the error
    verify(rich[0], times(0)).onError(any(LocalTransportException.class));
    // Wait for the notification
    if (!sync.await(TestingUtils.TESTING_DURATION().toMillis(), TimeUnit.MILLISECONDS)) {
        fail("Timed out after waiting for " + TestingUtils.TESTING_DURATION().toMillis() + " ms to be notified about the channel error.");
    }
    shutdown(serverAndClient);
}

Example 15 with InvocationOnMock

use of org.mockito.invocation.InvocationOnMock in project flink by apache.

the class NetworkEnvironmentTest method createSingleInputGateMock.

/**
	 * Helper to create a mock of a {@link SingleInputGate} for use by a {@link Task} inside
	 * {@link NetworkEnvironment#registerTask(Task)}.
	 *
	 * @param partitionType
	 * 		the consumed partition type
	 * @param channels
	 * 		the nummer of input channels
	 *
	 * @return mock with minimal functionality necessary by {@link NetworkEnvironment#registerTask(Task)}
	 */
private static SingleInputGate createSingleInputGateMock(final ResultPartitionType partitionType, final int channels) {
    SingleInputGate ig = mock(SingleInputGate.class);
    when(ig.getConsumedPartitionType()).thenReturn(partitionType);
    when(ig.getNumberOfInputChannels()).thenReturn(channels);
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(final InvocationOnMock invocation) throws Throwable {
            BufferPool bp = invocation.getArgumentAt(0, BufferPool.class);
            if (partitionType == ResultPartitionType.PIPELINED_BOUNDED) {
                assertEquals(channels * 2 + 8, bp.getMaxNumberOfMemorySegments());
            } else {
                assertEquals(Integer.MAX_VALUE, bp.getMaxNumberOfMemorySegments());
            }
            return null;
        }
    }).when(ig).setBufferPool(any(BufferPool.class));
    return ig;
}