Examples with MockEnvironmentBuilder org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder used on opensource projects

Example 11 with MockEnvironmentBuilder

use of org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder in project flink by apache.

the class StreamTaskTest method testExecuteMailboxActionsAfterLeavingInputProcessorMailboxLoop.

@Test
public void testExecuteMailboxActionsAfterLeavingInputProcessorMailboxLoop() throws Exception {
    OneShotLatch latch = new OneShotLatch();
    try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder().build()) {
        RunningTask<StreamTask<?, ?>> task = runTask(() -> new StreamTask<Object, StreamOperator<Object>>(mockEnvironment) {

            @Override
            protected void init() throws Exception {
            }

            @Override
            protected void processInput(MailboxDefaultAction.Controller controller) throws Exception {
                mailboxProcessor.getMailboxExecutor(0).execute(latch::trigger, "trigger");
                controller.suspendDefaultAction();
                mailboxProcessor.suspend();
            }
        });
        latch.await();
        task.waitForTaskCompletion(false);
    }
}

Example 12 with MockEnvironmentBuilder

use of org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder in project flink by apache.

the class StreamTaskTest method testFailingAsyncCheckpointRunnable.

/**
 * Tests that in case of a failing AsyncCheckpointRunnable all operator snapshot results are
 * cancelled and all non partitioned state handles are discarded.
 */
@Test
public void testFailingAsyncCheckpointRunnable() throws Exception {
    // mock the new state operator snapshots
    OperatorSnapshotFutures operatorSnapshotResult1 = mock(OperatorSnapshotFutures.class);
    OperatorSnapshotFutures operatorSnapshotResult2 = mock(OperatorSnapshotFutures.class);
    OperatorSnapshotFutures operatorSnapshotResult3 = mock(OperatorSnapshotFutures.class);
    RunnableFuture<SnapshotResult<OperatorStateHandle>> failingFuture = mock(RunnableFuture.class);
    when(failingFuture.get()).thenThrow(new ExecutionException(new Exception("Test exception")));
    when(operatorSnapshotResult3.getOperatorStateRawFuture()).thenReturn(failingFuture);
    try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder().build()) {
        RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(mockEnvironment, operatorChain(streamOperatorWithSnapshot(operatorSnapshotResult1), streamOperatorWithSnapshot(operatorSnapshotResult2), streamOperatorWithSnapshot(operatorSnapshotResult3))));
        MockStreamTask streamTask = task.streamTask;
        waitTaskIsRunning(streamTask, task.invocationFuture);
        mockEnvironment.setExpectedExternalFailureCause(Throwable.class);
        streamTask.triggerCheckpointAsync(new CheckpointMetaData(42L, 1L), CheckpointOptions.forCheckpointWithDefaultLocation()).get();
        // wait for the completion of the async task
        ExecutorService executor = streamTask.getAsyncOperationsThreadPool();
        executor.shutdown();
        if (!executor.awaitTermination(10000L, TimeUnit.MILLISECONDS)) {
            fail("Executor did not shut down within the given timeout. This indicates that the " + "checkpointing did not resume.");
        }
        assertTrue(mockEnvironment.getActualExternalFailureCause().isPresent());
        verify(operatorSnapshotResult1).cancel();
        verify(operatorSnapshotResult2).cancel();
        verify(operatorSnapshotResult3).cancel();
        streamTask.finishInput();
        task.waitForTaskCompletion(false);
    }
}

Example 13 with MockEnvironmentBuilder

use of org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder in project flink by apache.

the class StreamTaskTest method testThreadInvariants.

/**
 * Tests that some StreamTask methods are called only in the main task's thread. Currently, the
 * main task's thread is the thread that creates the task.
 */
@Test
public void testThreadInvariants() throws Throwable {
    Configuration taskConfiguration = new Configuration();
    StreamConfig streamConfig = new StreamConfig(taskConfiguration);
    streamConfig.setStreamOperator(new StreamMap<>(value -> value));
    streamConfig.setOperatorID(new OperatorID());
    try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder().setTaskConfiguration(taskConfiguration).build()) {
        ClassLoader taskClassLoader = new TestUserCodeClassLoader();
        RunningTask<ThreadInspectingTask> runningTask = runTask(() -> {
            Thread.currentThread().setContextClassLoader(taskClassLoader);
            return new ThreadInspectingTask(mockEnvironment);
        });
        runningTask.invocationFuture.get();
        assertThat(runningTask.streamTask.getTaskClassLoader(), is(sameInstance(taskClassLoader)));
    }
}

Example 14 with MockEnvironmentBuilder

use of org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder in project flink by apache.

the class StreamTaskTest method setupEnvironment.

private MockEnvironment setupEnvironment(boolean... outputAvailabilities) {
    final Configuration configuration = new Configuration();
    new MockStreamConfig(configuration, outputAvailabilities.length);
    final List<ResultPartitionWriter> writers = new ArrayList<>(outputAvailabilities.length);
    for (int i = 0; i < outputAvailabilities.length; i++) {
        writers.add(new AvailabilityTestResultPartitionWriter(outputAvailabilities[i]));
    }
    final MockEnvironment environment = new MockEnvironmentBuilder().setTaskConfiguration(configuration).build();
    environment.addOutputs(writers);
    return environment;
}

Example 15 with MockEnvironmentBuilder

use of org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder in project flink by apache.

the class StreamTaskTest method testAsyncCheckpointingConcurrentCloseAfterAcknowledge.

/**
 * FLINK-5667
 *
 * <p>Tests that a concurrent cancel operation does not discard the state handles of an
 * acknowledged checkpoint. The situation can only happen if the cancel call is executed after
 * Environment.acknowledgeCheckpoint() and before the CloseableRegistry.unregisterClosable()
 * call.
 */
@Test
public void testAsyncCheckpointingConcurrentCloseAfterAcknowledge() throws Exception {
    final OneShotLatch acknowledgeCheckpointLatch = new OneShotLatch();
    final OneShotLatch completeAcknowledge = new OneShotLatch();
    CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) {
            acknowledgeCheckpointLatch.trigger();
            // block here so that we can issue the concurrent cancel call
            while (true) {
                try {
                    // wait until we successfully await (no pun intended)
                    completeAcknowledge.await();
                    // when await() returns normally, we break out of the loop
                    break;
                } catch (InterruptedException e) {
                // survive interruptions that arise from thread pool
                // shutdown
                // production code cannot actually throw
                // InterruptedException from
                // checkpoint acknowledgement
                }
            }
            return null;
        }
    }).when(checkpointResponder).acknowledgeCheckpoint(any(JobID.class), any(ExecutionAttemptID.class), anyLong(), any(CheckpointMetrics.class), any(TaskStateSnapshot.class));
    TaskStateManager taskStateManager = new TaskStateManagerImpl(new JobID(1L, 2L), new ExecutionAttemptID(), mock(TaskLocalStateStoreImpl.class), new InMemoryStateChangelogStorage(), null, checkpointResponder);
    KeyedStateHandle managedKeyedStateHandle = mock(KeyedStateHandle.class);
    KeyedStateHandle rawKeyedStateHandle = mock(KeyedStateHandle.class);
    OperatorStateHandle managedOperatorStateHandle = mock(OperatorStreamStateHandle.class);
    OperatorStateHandle rawOperatorStateHandle = mock(OperatorStreamStateHandle.class);
    OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(DoneFuture.of(SnapshotResult.of(managedKeyedStateHandle)), DoneFuture.of(SnapshotResult.of(rawKeyedStateHandle)), DoneFuture.of(SnapshotResult.of(managedOperatorStateHandle)), DoneFuture.of(SnapshotResult.of(rawOperatorStateHandle)), DoneFuture.of(SnapshotResult.empty()), DoneFuture.of(SnapshotResult.empty()));
    try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder().setTaskName("mock-task").setTaskStateManager(taskStateManager).build()) {
        RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(mockEnvironment, operatorChain(streamOperatorWithSnapshot(operatorSnapshotResult))));
        MockStreamTask streamTask = task.streamTask;
        waitTaskIsRunning(streamTask, task.invocationFuture);
        final long checkpointId = 42L;
        streamTask.triggerCheckpointAsync(new CheckpointMetaData(checkpointId, 1L), CheckpointOptions.forCheckpointWithDefaultLocation());
        acknowledgeCheckpointLatch.await();
        ArgumentCaptor<TaskStateSnapshot> subtaskStateCaptor = ArgumentCaptor.forClass(TaskStateSnapshot.class);
        // check that the checkpoint has been completed
        verify(checkpointResponder).acknowledgeCheckpoint(any(JobID.class), any(ExecutionAttemptID.class), eq(checkpointId), any(CheckpointMetrics.class), subtaskStateCaptor.capture());
        TaskStateSnapshot subtaskStates = subtaskStateCaptor.getValue();
        OperatorSubtaskState subtaskState = subtaskStates.getSubtaskStateMappings().iterator().next().getValue();
        // check that the subtask state contains the expected state handles
        assertEquals(singleton(managedKeyedStateHandle), subtaskState.getManagedKeyedState());
        assertEquals(singleton(rawKeyedStateHandle), subtaskState.getRawKeyedState());
        assertEquals(singleton(managedOperatorStateHandle), subtaskState.getManagedOperatorState());
        assertEquals(singleton(rawOperatorStateHandle), subtaskState.getRawOperatorState());
        // check that the state handles have not been discarded
        verify(managedKeyedStateHandle, never()).discardState();
        verify(rawKeyedStateHandle, never()).discardState();
        verify(managedOperatorStateHandle, never()).discardState();
        verify(rawOperatorStateHandle, never()).discardState();
        streamTask.cancel();
        completeAcknowledge.trigger();
        // canceling the stream task after it has acknowledged the checkpoint should not discard
        // the state handles
        verify(managedKeyedStateHandle, never()).discardState();
        verify(rawKeyedStateHandle, never()).discardState();
        verify(managedOperatorStateHandle, never()).discardState();
        verify(rawOperatorStateHandle, never()).discardState();
        task.waitForTaskCompletion(true);
    }
}