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

Example 36 with MockEnvironment

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

the class SortingDataInputTest method simpleFixedLengthKeySorting.

@Test
public void simpleFixedLengthKeySorting() throws Exception {
    CollectingDataOutput<Integer> collectingDataOutput = new CollectingDataOutput<>();
    CollectionDataInput<Integer> input = new CollectionDataInput<>(Arrays.asList(new StreamRecord<>(1, 3), new StreamRecord<>(1, 1), new StreamRecord<>(2, 1), new StreamRecord<>(2, 3), new StreamRecord<>(1, 2), new StreamRecord<>(2, 2)));
    MockEnvironment environment = MockEnvironment.builder().build();
    SortingDataInput<Integer, Integer> sortingDataInput = new SortingDataInput<>(input, new IntSerializer(), new IntSerializer(), (KeySelector<Integer, Integer>) value -> value, environment.getMemoryManager(), environment.getIOManager(), true, 1.0, new Configuration(), new DummyInvokable(), new ExecutionConfig());
    DataInputStatus inputStatus;
    do {
        inputStatus = sortingDataInput.emitNext(collectingDataOutput);
    } while (inputStatus != DataInputStatus.END_OF_INPUT);
    assertThat(collectingDataOutput.events, equalTo(Arrays.asList(new StreamRecord<>(1, 1), new StreamRecord<>(1, 2), new StreamRecord<>(1, 3), new StreamRecord<>(2, 1), new StreamRecord<>(2, 2), new StreamRecord<>(2, 3))));
}

Example 37 with MockEnvironment

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

the class SortingDataInputTest method watermarkPropagation.

@Test
public void watermarkPropagation() throws Exception {
    CollectingDataOutput<Integer> collectingDataOutput = new CollectingDataOutput<>();
    CollectionDataInput<Integer> input = new CollectionDataInput<>(Arrays.asList(new StreamRecord<>(1, 3), new Watermark(1), new StreamRecord<>(1, 1), new Watermark(2), new StreamRecord<>(2, 1), new Watermark(3), new StreamRecord<>(2, 3), new Watermark(4), new StreamRecord<>(1, 2), new Watermark(5), new StreamRecord<>(2, 2), new Watermark(6)));
    MockEnvironment environment = MockEnvironment.builder().build();
    SortingDataInput<Integer, Integer> sortingDataInput = new SortingDataInput<>(input, new IntSerializer(), new IntSerializer(), (KeySelector<Integer, Integer>) value -> value, environment.getMemoryManager(), environment.getIOManager(), true, 1.0, new Configuration(), new DummyInvokable(), new ExecutionConfig());
    DataInputStatus inputStatus;
    do {
        inputStatus = sortingDataInput.emitNext(collectingDataOutput);
    } while (inputStatus != DataInputStatus.END_OF_INPUT);
    assertThat(collectingDataOutput.events, equalTo(Arrays.asList(new StreamRecord<>(1, 1), new StreamRecord<>(1, 2), new StreamRecord<>(1, 3), new StreamRecord<>(2, 1), new StreamRecord<>(2, 2), new StreamRecord<>(2, 3), new Watermark(6))));
}

Example 38 with MockEnvironment

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

the class StreamTaskTest method testProcessWithUnAvailableInput.

@Test
public void testProcessWithUnAvailableInput() throws Exception {
    final long sleepTimeOutsideMail = 42;
    final long sleepTimeInsideMail = 44;
    @Nullable WaitingThread waitingThread = null;
    try (final MockEnvironment environment = setupEnvironment(true, true)) {
        final UnAvailableTestInputProcessor inputProcessor = new UnAvailableTestInputProcessor();
        final StreamTask task = new MockStreamTaskBuilder(environment).setStreamInputProcessor(inputProcessor).build();
        TaskIOMetricGroup ioMetricGroup = task.getEnvironment().getMetricGroup().getIOMetricGroup();
        final MailboxExecutor executor = task.mailboxProcessor.getMainMailboxExecutor();
        final RunnableWithException completeFutureTask = () -> {
            inputProcessor.availabilityProvider.getUnavailableToResetAvailable().complete(null);
        };
        waitingThread = new WaitingThread(executor, completeFutureTask, sleepTimeInsideMail, sleepTimeOutsideMail, ioMetricGroup.getIdleTimeMsPerSecond());
        // Make sure WaitingThread is started after Task starts processing.
        executor.submit(waitingThread::start, "Start WaitingThread after Task starts processing input.");
        SystemClock clock = SystemClock.getInstance();
        long startTs = clock.absoluteTimeMillis();
        task.invoke();
        long totalDuration = clock.absoluteTimeMillis() - startTs;
        assertThat(ioMetricGroup.getIdleTimeMsPerSecond().getCount(), greaterThanOrEqualTo(sleepTimeOutsideMail));
        assertThat(ioMetricGroup.getIdleTimeMsPerSecond().getCount(), Matchers.lessThanOrEqualTo(totalDuration - sleepTimeInsideMail));
        assertThat(ioMetricGroup.getSoftBackPressuredTimePerSecond().getCount(), is(0L));
        assertThat(ioMetricGroup.getHardBackPressuredTimePerSecond().getCount(), is(0L));
    } finally {
        if (waitingThread != null) {
            waitingThread.join();
        }
    }
}

Example 39 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment 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);
    }
}

Example 40 with MockEnvironment

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

the class StreamTaskTest method testEmptySubtaskStateLeadsToStatelessAcknowledgment.

/**
 * FLINK-5985
 *
 * <p>This test ensures that empty snapshots (no op/keyed stated whatsoever) will be reported as
 * stateless tasks. This happens by translating an empty {@link SubtaskState} into reporting
 * 'null' to #acknowledgeCheckpoint.
 */
@Test
public void testEmptySubtaskStateLeadsToStatelessAcknowledgment() throws Exception {
    // latch blocks until the async checkpoint thread acknowledges
    final OneShotLatch checkpointCompletedLatch = new OneShotLatch();
    final List<SubtaskState> checkpointResult = new ArrayList<>(1);
    CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            SubtaskState subtaskState = invocation.getArgument(4);
            checkpointResult.add(subtaskState);
            checkpointCompletedLatch.trigger();
            return null;
        }
    }).when(checkpointResponder).acknowledgeCheckpoint(any(JobID.class), any(ExecutionAttemptID.class), anyLong(), any(CheckpointMetrics.class), nullable(TaskStateSnapshot.class));
    TaskStateManager taskStateManager = new TaskStateManagerImpl(new JobID(1L, 2L), new ExecutionAttemptID(), mock(TaskLocalStateStoreImpl.class), new InMemoryStateChangelogStorage(), null, checkpointResponder);
    // mock the operator with empty snapshot result (all state handles are null)
    OneInputStreamOperator<String, String> statelessOperator = streamOperatorWithSnapshot(new OperatorSnapshotFutures());
    try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder().setTaskStateManager(taskStateManager).build()) {
        RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(mockEnvironment, operatorChain(statelessOperator)));
        waitTaskIsRunning(task.streamTask, task.invocationFuture);
        task.streamTask.triggerCheckpointAsync(new CheckpointMetaData(42L, 1L), CheckpointOptions.forCheckpointWithDefaultLocation());
        checkpointCompletedLatch.await(30, TimeUnit.SECONDS);
        // ensure that 'null' was acknowledged as subtask state
        Assert.assertNull(checkpointResult.get(0));
        task.streamTask.cancel();
        task.waitForTaskCompletion(true);
    }
}