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

Example 16 with InvocationOnMock

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

the class RecordWriterTest method createBufferProvider.

private BufferProvider createBufferProvider(final int bufferSize) throws IOException, InterruptedException {
    BufferProvider bufferProvider = mock(BufferProvider.class);
    when(bufferProvider.requestBufferBlocking()).thenAnswer(new Answer<Buffer>() {

        @Override
        public Buffer answer(InvocationOnMock invocationOnMock) throws Throwable {
            MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(bufferSize);
            Buffer buffer = new Buffer(segment, DiscardingRecycler.INSTANCE);
            return buffer;
        }
    });
    return bufferProvider;
}

Example 17 with InvocationOnMock

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

the class SpillableSubpartitionTest method testConcurrentFinishAndReleaseMemory.

/**
	 * Tests a fix for FLINK-2384.
	 *
	 * @see <a href="https://issues.apache.org/jira/browse/FLINK-2384">FLINK-2384</a>
	 */
@Test
public void testConcurrentFinishAndReleaseMemory() throws Exception {
    // Latches to blocking
    final CountDownLatch doneLatch = new CountDownLatch(1);
    final CountDownLatch blockLatch = new CountDownLatch(1);
    // Blocking spill writer (blocks on the close call)
    AsynchronousBufferFileWriter spillWriter = mock(AsynchronousBufferFileWriter.class);
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            blockLatch.countDown();
            doneLatch.await();
            return null;
        }
    }).when(spillWriter).close();
    // Mock I/O manager returning the blocking spill writer
    IOManager ioManager = mock(IOManager.class);
    when(ioManager.createBufferFileWriter(any(FileIOChannel.ID.class))).thenReturn(spillWriter);
    // The partition
    final SpillableSubpartition partition = new SpillableSubpartition(0, mock(ResultPartition.class), ioManager);
    // Spill the partition initially (creates the spill writer)
    assertEquals(0, partition.releaseMemory());
    ExecutorService executor = Executors.newSingleThreadExecutor();
    // Finish the partition (this blocks because of the mock blocking writer)
    Future<Void> blockingFinish = executor.submit(new Callable<Void>() {

        @Override
        public Void call() throws Exception {
            partition.finish();
            return null;
        }
    });
    // Ensure that the blocking call has been made
    blockLatch.await();
    // This call needs to go through. FLINK-2384 discovered a bug, in
    // which the finish call was holding a lock, which was leading to a
    // deadlock when another operation on the partition was happening.
    partition.releaseMemory();
    // Check that the finish call succeeded w/o problems as well to avoid
    // false test successes.
    doneLatch.countDown();
    blockingFinish.get();
}

Example 18 with InvocationOnMock

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

the class ZooKeeperStateHandleStoreITCase method testRemoveWithCallback.

/**
	 * Tests that state handles are correctly removed with a callback.
	 */
@Test
public void testRemoveWithCallback() throws Exception {
    // Setup
    LongStateStorage stateHandleProvider = new LongStateStorage();
    ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(ZooKeeper.getClient(), stateHandleProvider, Executors.directExecutor());
    // Config
    final String pathInZooKeeper = "/testRemoveWithCallback";
    final Long state = 27255442L;
    store.add(pathInZooKeeper, state);
    final CountDownLatch sync = new CountDownLatch(1);
    BackgroundCallback callback = mock(BackgroundCallback.class);
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            sync.countDown();
            return null;
        }
    }).when(callback).processResult(eq(ZooKeeper.getClient()), any(CuratorEvent.class));
    // Test
    store.remove(pathInZooKeeper, callback);
    // Verify discarded and callback called
    assertEquals(0, ZooKeeper.getClient().getChildren().forPath("/").size());
    sync.await();
    verify(callback, times(1)).processResult(eq(ZooKeeper.getClient()), any(CuratorEvent.class));
}

Example 19 with InvocationOnMock

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

the class StreamTestSingleInputGate method setupInputChannels.

@SuppressWarnings("unchecked")
private void setupInputChannels() throws IOException, InterruptedException {
    for (int i = 0; i < numInputChannels; i++) {
        final int channelIndex = i;
        final RecordSerializer<SerializationDelegate<Object>> recordSerializer = new SpanningRecordSerializer<SerializationDelegate<Object>>();
        final SerializationDelegate<Object> delegate = (SerializationDelegate<Object>) (SerializationDelegate<?>) new SerializationDelegate<StreamElement>(new StreamElementSerializer<T>(serializer));
        inputQueues[channelIndex] = new ConcurrentLinkedQueue<InputValue<Object>>();
        inputChannels[channelIndex] = new TestInputChannel(inputGate, i);
        final Answer<BufferAndAvailability> answer = new Answer<BufferAndAvailability>() {

            @Override
            public BufferAndAvailability answer(InvocationOnMock invocationOnMock) throws Throwable {
                InputValue<Object> input = inputQueues[channelIndex].poll();
                if (input != null && input.isStreamEnd()) {
                    when(inputChannels[channelIndex].getInputChannel().isReleased()).thenReturn(true);
                    return new BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE), false);
                } else if (input != null && input.isStreamRecord()) {
                    Object inputElement = input.getStreamRecord();
                    final Buffer buffer = new Buffer(MemorySegmentFactory.allocateUnpooledSegment(bufferSize), mock(BufferRecycler.class));
                    recordSerializer.setNextBuffer(buffer);
                    delegate.setInstance(inputElement);
                    recordSerializer.addRecord(delegate);
                    // Call getCurrentBuffer to ensure size is set
                    return new BufferAndAvailability(recordSerializer.getCurrentBuffer(), false);
                } else if (input != null && input.isEvent()) {
                    AbstractEvent event = input.getEvent();
                    return new BufferAndAvailability(EventSerializer.toBuffer(event), false);
                } else {
                    synchronized (inputQueues[channelIndex]) {
                        inputQueues[channelIndex].wait();
                        return answer(invocationOnMock);
                    }
                }
            }
        };
        when(inputChannels[channelIndex].getInputChannel().getNextBuffer()).thenAnswer(answer);
        inputGate.setInputChannel(new IntermediateResultPartitionID(), inputChannels[channelIndex].getInputChannel());
    }
}

Example 20 with InvocationOnMock

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

the class StreamingRuntimeContextTest method createListPlainMockOp.

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception {
    AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class);
    ExecutionConfig config = new ExecutionConfig();
    KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);
    DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);
    when(operatorMock.getExecutionConfig()).thenReturn(config);
    doAnswer(new Answer<ListState<String>>() {

        @Override
        public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
            ListStateDescriptor<String> descr = (ListStateDescriptor<String>) invocationOnMock.getArguments()[2];
            AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(new DummyEnvironment("test_task", 1, 0), new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
            backend.setCurrentKey(0);
            return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
        }
    }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));
    when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
    return operatorMock;
}