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

Example 26 with MockEnvironment

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

the class SubtaskCheckpointCoordinatorTest method testNotifyCheckpointAbortedDuringAsyncPhase.

@Test
public void testNotifyCheckpointAbortedDuringAsyncPhase() throws Exception {
    MockEnvironment mockEnvironment = MockEnvironment.builder().build();
    try (SubtaskCheckpointCoordinatorImpl subtaskCheckpointCoordinator = (SubtaskCheckpointCoordinatorImpl) new MockSubtaskCheckpointCoordinatorBuilder().setEnvironment(mockEnvironment).setExecutor(Executors.newSingleThreadExecutor()).setUnalignedCheckpointEnabled(true).build()) {
        final BlockingRunnableFuture rawKeyedStateHandleFuture = new BlockingRunnableFuture();
        OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(DoneFuture.of(SnapshotResult.empty()), rawKeyedStateHandleFuture, DoneFuture.of(SnapshotResult.empty()), DoneFuture.of(SnapshotResult.empty()), DoneFuture.of(SnapshotResult.empty()), DoneFuture.of(SnapshotResult.empty()));
        final OperatorChain<String, AbstractStreamOperator<String>> operatorChain = operatorChain(new CheckpointOperator(operatorSnapshotResult));
        long checkpointId = 42L;
        subtaskCheckpointCoordinator.getChannelStateWriter().start(checkpointId, CheckpointOptions.forCheckpointWithDefaultLocation());
        subtaskCheckpointCoordinator.checkpointState(new CheckpointMetaData(checkpointId, System.currentTimeMillis()), CheckpointOptions.forCheckpointWithDefaultLocation(), new CheckpointMetricsBuilder(), operatorChain, false, () -> false);
        rawKeyedStateHandleFuture.awaitRun();
        assertEquals(1, subtaskCheckpointCoordinator.getAsyncCheckpointRunnableSize());
        assertFalse(rawKeyedStateHandleFuture.isCancelled());
        subtaskCheckpointCoordinator.notifyCheckpointAborted(checkpointId, operatorChain, () -> true);
        assertTrue(rawKeyedStateHandleFuture.isCancelled());
        assertEquals(0, subtaskCheckpointCoordinator.getAsyncCheckpointRunnableSize());
    }
}

Example 27 with MockEnvironment

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

the class SubtaskCheckpointCoordinatorTest method testNotifyCheckpointAbortedAfterAsyncPhase.

@Test
public void testNotifyCheckpointAbortedAfterAsyncPhase() throws Exception {
    TestTaskStateManager stateManager = new TestTaskStateManager();
    MockEnvironment mockEnvironment = MockEnvironment.builder().setTaskStateManager(stateManager).build();
    try (SubtaskCheckpointCoordinatorImpl subtaskCheckpointCoordinator = (SubtaskCheckpointCoordinatorImpl) new MockSubtaskCheckpointCoordinatorBuilder().setEnvironment(mockEnvironment).build()) {
        final OperatorChain<?, ?> operatorChain = getOperatorChain(mockEnvironment);
        long checkpointId = 42L;
        subtaskCheckpointCoordinator.checkpointState(new CheckpointMetaData(checkpointId, System.currentTimeMillis()), CheckpointOptions.forCheckpointWithDefaultLocation(), new CheckpointMetricsBuilder(), operatorChain, false, () -> false);
        subtaskCheckpointCoordinator.notifyCheckpointAborted(checkpointId, operatorChain, () -> true);
        assertEquals(0, subtaskCheckpointCoordinator.getAbortedCheckpointSize());
        assertEquals(checkpointId, stateManager.getNotifiedAbortedCheckpointId());
    }
}

Example 28 with MockEnvironment

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

the class SubtaskCheckpointCoordinatorTest method testSavepointNotResultingInPriorityEvents.

@Test
public void testSavepointNotResultingInPriorityEvents() throws Exception {
    MockEnvironment mockEnvironment = MockEnvironment.builder().build();
    try (SubtaskCheckpointCoordinator coordinator = new MockSubtaskCheckpointCoordinatorBuilder().setUnalignedCheckpointEnabled(true).setEnvironment(mockEnvironment).build()) {
        AtomicReference<Boolean> broadcastedPriorityEvent = new AtomicReference<>(null);
        final OperatorChain<?, ?> operatorChain = new RegularOperatorChain(new MockStreamTaskBuilder(mockEnvironment).build(), new NonRecordWriter<>()) {

            @Override
            public void broadcastEvent(AbstractEvent event, boolean isPriorityEvent) throws IOException {
                super.broadcastEvent(event, isPriorityEvent);
                broadcastedPriorityEvent.set(isPriorityEvent);
            }
        };
        coordinator.checkpointState(new CheckpointMetaData(0, 0), new CheckpointOptions(SavepointType.savepoint(SavepointFormatType.CANONICAL), CheckpointStorageLocationReference.getDefault()), new CheckpointMetricsBuilder(), operatorChain, false, () -> true);
        assertEquals(false, broadcastedPriorityEvent.get());
    }
}

Example 29 with MockEnvironment

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

the class SubtaskCheckpointCoordinatorTest method testNotifyCheckpointAbortedManyTimes.

@Test
public void testNotifyCheckpointAbortedManyTimes() throws Exception {
    MockEnvironment mockEnvironment = MockEnvironment.builder().build();
    int maxRecordAbortedCheckpoints = 256;
    try (SubtaskCheckpointCoordinatorImpl subtaskCheckpointCoordinator = (SubtaskCheckpointCoordinatorImpl) new MockSubtaskCheckpointCoordinatorBuilder().setEnvironment(mockEnvironment).setMaxRecordAbortedCheckpoints(maxRecordAbortedCheckpoints).build()) {
        final OperatorChain<?, ?> operatorChain = getOperatorChain(mockEnvironment);
        long notifyAbortedTimes = maxRecordAbortedCheckpoints + 42;
        for (int i = 1; i < notifyAbortedTimes; i++) {
            subtaskCheckpointCoordinator.notifyCheckpointAborted(i, operatorChain, () -> true);
            assertEquals(Math.min(maxRecordAbortedCheckpoints, i), subtaskCheckpointCoordinator.getAbortedCheckpointSize());
        }
    }
}

Example 30 with MockEnvironment

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

the class SubtaskCheckpointCoordinatorTest method testForceAlignedCheckpointResultingInPriorityEvents.

@Test
public void testForceAlignedCheckpointResultingInPriorityEvents() throws Exception {
    final long checkpointId = 42L;
    MockEnvironment mockEnvironment = MockEnvironment.builder().build();
    try (SubtaskCheckpointCoordinator coordinator = new MockSubtaskCheckpointCoordinatorBuilder().setUnalignedCheckpointEnabled(true).setEnvironment(mockEnvironment).build()) {
        AtomicReference<Boolean> broadcastedPriorityEvent = new AtomicReference<>(null);
        final OperatorChain<?, ?> operatorChain = new RegularOperatorChain(new MockStreamTaskBuilder(mockEnvironment).build(), new NonRecordWriter<>()) {

            @Override
            public void broadcastEvent(AbstractEvent event, boolean isPriorityEvent) throws IOException {
                super.broadcastEvent(event, isPriorityEvent);
                broadcastedPriorityEvent.set(isPriorityEvent);
                // test if we can write output data
                coordinator.getChannelStateWriter().addOutputData(checkpointId, new ResultSubpartitionInfo(0, 0), 0, BufferBuilderTestUtils.buildSomeBuffer(500));
            }
        };
        CheckpointOptions forcedAlignedOptions = CheckpointOptions.unaligned(CheckpointType.CHECKPOINT, CheckpointStorageLocationReference.getDefault()).withUnalignedUnsupported();
        coordinator.checkpointState(new CheckpointMetaData(checkpointId, 0), forcedAlignedOptions, new CheckpointMetricsBuilder(), operatorChain, false, () -> true);
        assertEquals(true, broadcastedPriorityEvent.get());
    }
}