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

Example 1 with MockEnvironmentBuilder

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

the class SourceOperatorEventTimeTest method createTestOperator.

// ------------------------------------------------------------------------
// test setup helpers
// ------------------------------------------------------------------------
private static <T> SourceOperator<T, MockSourceSplit> createTestOperator(SourceReader<T, MockSourceSplit> reader, WatermarkStrategy<T> watermarkStrategy, ProcessingTimeService timeService, boolean emitProgressiveWatermarks) throws Exception {
    final OperatorStateStore operatorStateStore = new MemoryStateBackend().createOperatorStateBackend(new MockEnvironmentBuilder().build(), "test-operator", Collections.emptyList(), new CloseableRegistry());
    final StateInitializationContext stateContext = new StateInitializationContextImpl(null, operatorStateStore, null, null, null);
    final SourceOperator<T, MockSourceSplit> sourceOperator = new TestingSourceOperator<>(reader, watermarkStrategy, timeService, emitProgressiveWatermarks);
    sourceOperator.setup(new SourceOperatorStreamTask<Integer>(new StreamMockEnvironment(new Configuration(), new Configuration(), new ExecutionConfig(), 1L, new MockInputSplitProvider(), 1, new TestTaskStateManager())), new MockStreamConfig(new Configuration(), 1), new MockOutput<>(new ArrayList<>()));
    sourceOperator.initializeState(stateContext);
    sourceOperator.open();
    return sourceOperator;
}

Example 2 with MockEnvironmentBuilder

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

the class SourceOperatorTestContext method createOperatorStateStore.

private OperatorStateStore createOperatorStateStore() throws Exception {
    MockEnvironment env = new MockEnvironmentBuilder().build();
    final AbstractStateBackend abstractStateBackend = new MemoryStateBackend();
    CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
    return abstractStateBackend.createOperatorStateBackend(env, "test-operator", Collections.emptyList(), cancelStreamRegistry);
}

Example 3 with MockEnvironmentBuilder

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

the class StreamOperatorStateHandlerTest method testFailingBackendSnapshotMethod.

/**
 * Tests that a failing snapshot method call to the keyed state backend will trigger the closing
 * of the StateSnapshotContextSynchronousImpl and the cancellation of the
 * OperatorSnapshotResult. The latter is supposed to also cancel all assigned futures.
 */
@Test
public void testFailingBackendSnapshotMethod() throws Exception {
    final long checkpointId = 42L;
    final long timestamp = 1L;
    try (CloseableRegistry closeableRegistry = new CloseableRegistry()) {
        RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateManagedFuture = new CancelableFuture<>();
        RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateRawFuture = new CancelableFuture<>();
        RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateManagedFuture = new CancelableFuture<>();
        RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateRawFuture = new CancelableFuture<>();
        RunnableFuture<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>> inputChannelStateFuture = new CancelableFuture<>();
        RunnableFuture<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>> resultSubpartitionStateFuture = new CancelableFuture<>();
        OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(keyedStateManagedFuture, keyedStateRawFuture, operatorStateManagedFuture, operatorStateRawFuture, inputChannelStateFuture, resultSubpartitionStateFuture);
        StateSnapshotContextSynchronousImpl context = new TestStateSnapshotContextSynchronousImpl(checkpointId, timestamp, closeableRegistry);
        context.getRawKeyedOperatorStateOutput();
        context.getRawOperatorStateOutput();
        StreamTaskStateInitializerImpl stateInitializer = new StreamTaskStateInitializerImpl(new MockEnvironmentBuilder().build(), new MemoryStateBackend());
        StreamOperatorStateContext stateContext = stateInitializer.streamOperatorStateContext(new OperatorID(), "whatever", new TestProcessingTimeService(), new UnUsedKeyContext(), IntSerializer.INSTANCE, closeableRegistry, new InterceptingOperatorMetricGroup(), 1.0, false);
        StreamOperatorStateHandler stateHandler = new StreamOperatorStateHandler(stateContext, new ExecutionConfig(), closeableRegistry);
        final String keyedStateField = "keyedStateField";
        final String operatorStateField = "operatorStateField";
        CheckpointedStreamOperator checkpointedStreamOperator = new CheckpointedStreamOperator() {

            @Override
            public void initializeState(StateInitializationContext context) throws Exception {
                context.getKeyedStateStore().getState(new ValueStateDescriptor<>(keyedStateField, LongSerializer.INSTANCE)).update(42L);
                context.getOperatorStateStore().getListState(new ListStateDescriptor<>(operatorStateField, LongSerializer.INSTANCE)).add(42L);
            }

            @Override
            public void snapshotState(StateSnapshotContext context) throws Exception {
                throw new ExpectedTestException();
            }
        };
        stateHandler.setCurrentKey("44");
        stateHandler.initializeOperatorState(checkpointedStreamOperator);
        assertThat(stateContext.operatorStateBackend().getRegisteredStateNames(), is(not(empty())));
        assertThat(((AbstractKeyedStateBackend<?>) stateContext.keyedStateBackend()).numKeyValueStatesByName(), equalTo(1));
        try {
            stateHandler.snapshotState(checkpointedStreamOperator, Optional.of(stateContext.internalTimerServiceManager()), "42", 42, 42, CheckpointOptions.forCheckpointWithDefaultLocation(), new MemCheckpointStreamFactory(1024), operatorSnapshotResult, context, false);
            fail("Exception expected.");
        } catch (CheckpointException e) {
            // as CheckpointException is wrapping the cause with SerializedThrowable
            if (!ExceptionUtils.findThrowableWithMessage(e, ExpectedTestException.MESSAGE).isPresent()) {
                throw e;
            }
        }
        assertTrue(keyedStateManagedFuture.isCancelled());
        assertTrue(keyedStateRawFuture.isCancelled());
        assertTrue(context.getKeyedStateStreamFuture().isCancelled());
        assertTrue(operatorStateManagedFuture.isCancelled());
        assertTrue(operatorStateRawFuture.isCancelled());
        assertTrue(context.getOperatorStateStreamFuture().isCancelled());
        assertTrue(inputChannelStateFuture.isCancelled());
        assertTrue(resultSubpartitionStateFuture.isCancelled());
        stateHandler.dispose();
        assertThat(stateContext.operatorStateBackend().getRegisteredBroadcastStateNames(), is(empty()));
        assertThat(stateContext.operatorStateBackend().getRegisteredStateNames(), is(empty()));
        assertThat(((AbstractKeyedStateBackend<?>) stateContext.keyedStateBackend()).numKeyValueStatesByName(), equalTo(0));
    }
}

Example 4 with MockEnvironmentBuilder

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

the class BufferDataOverWindowOperatorTest method test.

private void test(OverWindowFrame[] frames, GenericRowData[] expect) throws Exception {
    MockEnvironment env = new MockEnvironmentBuilder().setIOManager(ioManager).setMemoryManager(memoryManager).build();
    StreamTask<Object, StreamOperator<Object>> task = new StreamTask<Object, StreamOperator<Object>>(env) {

        @Override
        protected void init() {
        }
    };
    operator = new BufferDataOverWindowOperator(frames, comparator, true) {

        {
            output = new NonBufferOverWindowOperatorTest.ConsumerOutput(new Consumer<RowData>() {

                @Override
                public void accept(RowData r) {
                    collect.add(GenericRowData.of(r.getInt(0), r.getLong(1), r.getLong(2), r.getLong(3), r.getLong(4)));
                }
            });
        }

        @Override
        public ClassLoader getUserCodeClassloader() {
            return Thread.currentThread().getContextClassLoader();
        }

        @Override
        public StreamConfig getOperatorConfig() {
            StreamConfig conf = mock(StreamConfig.class);
            when(conf.<RowData>getTypeSerializerIn1(getUserCodeClassloader())).thenReturn(inputSer);
            when(conf.getManagedMemoryFractionOperatorUseCaseOfSlot(eq(ManagedMemoryUseCase.OPERATOR), any(Configuration.class), any(ClassLoader.class))).thenReturn(0.99);
            return conf;
        }

        @Override
        public StreamTask<?, ?> getContainingTask() {
            return task;
        }

        @Override
        public StreamingRuntimeContext getRuntimeContext() {
            return mock(StreamingRuntimeContext.class);
        }
    };
    operator.setProcessingTimeService(new TestProcessingTimeService());
    operator.open();
    addRow(0, 1L, 4L);
    /* 1 **/
    addRow(0, 1L, 1L);
    /* 2 **/
    addRow(0, 1L, 1L);
    /* 3 **/
    addRow(0, 1L, 1L);
    /* 4 **/
    addRow(1, 5L, 2L);
    /* 5 **/
    addRow(2, 5L, 4L);
    /* 6 **/
    addRow(2, 6L, 2L);
    /* 7 **/
    addRow(2, 6L, 2L);
    /* 8 **/
    addRow(2, 6L, 2L);
    /* 9 **/
    operator.endInput();
    GenericRowData[] outputs = this.collect.toArray(new GenericRowData[0]);
    Assert.assertArrayEquals(expect, outputs);
    operator.close();
}

Example 5 with MockEnvironmentBuilder

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

the class PojoSerializerUpgradeTest method runOperator.

private OperatorSubtaskState runOperator(Configuration taskConfiguration, ExecutionConfig executionConfig, OneInputStreamOperator<Long, Long> operator, KeySelector<Long, Long> keySelector, boolean isKeyedState, StateBackend stateBackend, ClassLoader classLoader, OperatorSubtaskState operatorSubtaskState, Iterable<Long> input) throws Exception {
    try (final MockEnvironment environment = new MockEnvironmentBuilder().setTaskName("test task").setManagedMemorySize(32 * 1024).setInputSplitProvider(new MockInputSplitProvider()).setBufferSize(256).setTaskConfiguration(taskConfiguration).setExecutionConfig(executionConfig).setMaxParallelism(16).setUserCodeClassLoader(classLoader).build()) {
        OneInputStreamOperatorTestHarness<Long, Long> harness = null;
        try {
            if (isKeyedState) {
                harness = new KeyedOneInputStreamOperatorTestHarness<>(operator, keySelector, BasicTypeInfo.LONG_TYPE_INFO, environment);
            } else {
                harness = new OneInputStreamOperatorTestHarness<>(operator, LongSerializer.INSTANCE, environment);
            }
            harness.setStateBackend(stateBackend);
            harness.setup();
            harness.initializeState(operatorSubtaskState);
            harness.open();
            long timestamp = 0L;
            for (Long value : input) {
                harness.processElement(value, timestamp++);
            }
            long checkpointId = 1L;
            long checkpointTimestamp = timestamp + 1L;
            return harness.snapshot(checkpointId, checkpointTimestamp);
        } finally {
            IOUtils.closeQuietly(harness);
        }
    }
}