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

Example 6 with DefaultKeyedStateStore

use of org.apache.flink.runtime.state.DefaultKeyedStateStore 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;
}