Examples with KeyedStateBackend org.apache.flink.runtime.state.KeyedStateBackend used on opensource projects

Example 6 with KeyedStateBackend

use of org.apache.flink.runtime.state.KeyedStateBackend in project flink by apache.

the class StreamingRuntimeContextTest method createMapPlainMockOp.

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createMapPlainMockOp() 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<MapState<Integer, String>>() {

        @Override
        public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable {
            MapStateDescriptor<Integer, String> descr = (MapStateDescriptor<Integer, 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()), TtlTimeProvider.DEFAULT, new UnregisteredMetricsGroup(), Collections.emptyList(), new CloseableRegistry());
            backend.setCurrentKey(0);
            return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
        }
    }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class));
    when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore);
    when(operatorMock.getOperatorID()).thenReturn(new OperatorID());
    when(operatorMock.getProcessingTimeService()).thenReturn(new TestProcessingTimeService());
    return operatorMock;
}

Example 7 with KeyedStateBackend

use of org.apache.flink.runtime.state.KeyedStateBackend in project flink by apache.

the class StreamingRuntimeContextTest method createDescriptorCapturingMockOp.

@SuppressWarnings("unchecked")
private static AbstractStreamOperator<?> createDescriptorCapturingMockOp(final AtomicReference<Object> ref, final ExecutionConfig config, Environment environment) throws Exception {
    AbstractStreamOperator<?> operator = new AbstractStreamOperator<Object>() {

        @Override
        public void setup(StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<Object>> output) {
            super.setup(containingTask, config, output);
        }
    };
    StreamConfig streamConfig = new StreamConfig(new Configuration());
    streamConfig.setOperatorID(new OperatorID());
    operator.setup(new MockStreamTaskBuilder(environment).setExecutionConfig(config).build(), streamConfig, new CollectorOutput<>(new ArrayList<>()));
    StreamTaskStateInitializer streamTaskStateManager = new StreamTaskStateInitializerImpl(environment, new MemoryStateBackend());
    KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class);
    DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            ref.set(invocationOnMock.getArguments()[2]);
            return null;
        }
    }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(StateDescriptor.class));
    operator.initializeState(streamTaskStateManager);
    operator.getRuntimeContext().setKeyedStateStore(keyedStateStore);
    return operator;
}