Examples with StreamingRuntimeContext org.apache.flink.streaming.api.operators.StreamingRuntimeContext used on opensource projects

Example 6 with StreamingRuntimeContext

use of org.apache.flink.streaming.api.operators.StreamingRuntimeContext in project flink by apache.

the class SpoutWrapperTest method testRunPrepare.

@SuppressWarnings({ "rawtypes", "unchecked" })
@Test
public void testRunPrepare() throws Exception {
    final StormConfig stormConfig = new StormConfig();
    stormConfig.put(this.r.nextInt(), this.r.nextInt());
    final Configuration flinkConfig = new Configuration();
    flinkConfig.setInteger("testKey", this.r.nextInt());
    final ExecutionConfig taskConfig = mock(ExecutionConfig.class);
    when(taskConfig.getGlobalJobParameters()).thenReturn(null).thenReturn(stormConfig).thenReturn(flinkConfig);
    final StreamingRuntimeContext taskContext = mock(StreamingRuntimeContext.class);
    when(taskContext.getExecutionConfig()).thenReturn(taskConfig);
    when(taskContext.getTaskName()).thenReturn("name");
    final IRichSpout spout = mock(IRichSpout.class);
    SpoutWrapper spoutWrapper = new SpoutWrapper(spout);
    spoutWrapper.setRuntimeContext(taskContext);
    spoutWrapper.cancel();
    // test without configuration
    spoutWrapper.run(mock(SourceContext.class));
    verify(spout).open(any(Map.class), any(TopologyContext.class), any(SpoutOutputCollector.class));
    // test with StormConfig
    spoutWrapper.run(mock(SourceContext.class));
    verify(spout).open(eq(stormConfig), any(TopologyContext.class), any(SpoutOutputCollector.class));
    // test with Configuration
    final TestDummySpout testSpout = new TestDummySpout();
    spoutWrapper = new SpoutWrapper(testSpout);
    spoutWrapper.setRuntimeContext(taskContext);
    spoutWrapper.cancel();
    spoutWrapper.run(mock(SourceContext.class));
    for (Entry<String, String> entry : flinkConfig.toMap().entrySet()) {
        Assert.assertEquals(entry.getValue(), testSpout.config.get(entry.getKey()));
    }
}

Example 7 with StreamingRuntimeContext

use of org.apache.flink.streaming.api.operators.StreamingRuntimeContext in project flink by apache.

the class SpoutWrapperTest method testRunExecuteFinite.

@Test
public void testRunExecuteFinite() throws Exception {
    final int numberOfCalls = this.r.nextInt(50);
    final LinkedList<Tuple1<Integer>> expectedResult = new LinkedList<Tuple1<Integer>>();
    for (int i = numberOfCalls - 1; i >= 0; --i) {
        expectedResult.add(new Tuple1<Integer>(new Integer(i)));
    }
    final StreamingRuntimeContext taskContext = mock(StreamingRuntimeContext.class);
    when(taskContext.getExecutionConfig()).thenReturn(mock(ExecutionConfig.class));
    when(taskContext.getTaskName()).thenReturn("name");
    final FiniteTestSpout spout = new FiniteTestSpout(numberOfCalls);
    final SpoutWrapper<Tuple1<Integer>> spoutWrapper = new SpoutWrapper<Tuple1<Integer>>(spout, -1);
    spoutWrapper.setRuntimeContext(taskContext);
    final TestContext collector = new TestContext();
    spoutWrapper.run(collector);
    Assert.assertEquals(expectedResult, collector.result);
}

Example 8 with StreamingRuntimeContext

use of org.apache.flink.streaming.api.operators.StreamingRuntimeContext in project flink by apache.

the class SpoutWrapperTest method testCancel.

@Test
public void testCancel() throws Exception {
    final int numberOfCalls = 5 + this.r.nextInt(5);
    final StreamingRuntimeContext taskContext = mock(StreamingRuntimeContext.class);
    when(taskContext.getExecutionConfig()).thenReturn(mock(ExecutionConfig.class));
    when(taskContext.getTaskName()).thenReturn("name");
    final IRichSpout spout = new FiniteTestSpout(numberOfCalls);
    final SpoutWrapper<Tuple1<Integer>> spoutWrapper = new SpoutWrapper<Tuple1<Integer>>(spout);
    spoutWrapper.setRuntimeContext(taskContext);
    spoutWrapper.cancel();
    final TestContext collector = new TestContext();
    spoutWrapper.run(collector);
    Assert.assertEquals(new LinkedList<Tuple1<Integer>>(), collector.result);
}

Example 9 with StreamingRuntimeContext

use of org.apache.flink.streaming.api.operators.StreamingRuntimeContext in project flink by apache.

the class SourceFunctionUtil method runSourceFunction.

public static <T extends Serializable> List<T> runSourceFunction(SourceFunction<T> sourceFunction) throws Exception {
    final List<T> outputs = new ArrayList<T>();
    if (sourceFunction instanceof RichFunction) {
        AbstractStreamOperator<?> operator = mock(AbstractStreamOperator.class);
        when(operator.getExecutionConfig()).thenReturn(new ExecutionConfig());
        RuntimeContext runtimeContext = new StreamingRuntimeContext(operator, new MockEnvironment("MockTask", 3 * 1024 * 1024, new MockInputSplitProvider(), 1024), new HashMap<String, Accumulator<?, ?>>());
        ((RichFunction) sourceFunction).setRuntimeContext(runtimeContext);
        ((RichFunction) sourceFunction).open(new Configuration());
    }
    try {
        SourceFunction.SourceContext<T> ctx = new CollectingSourceContext<T>(new Object(), outputs);
        sourceFunction.run(ctx);
    } catch (Exception e) {
        throw new RuntimeException("Cannot invoke source.", e);
    }
    return outputs;
}

Example 10 with StreamingRuntimeContext

use of org.apache.flink.streaming.api.operators.StreamingRuntimeContext in project flink by apache.

the class RMQSource method open.

@Override
public void open(Configuration config) throws Exception {
    super.open(config);
    try {
        connection = setupConnection();
        channel = setupChannel(connection);
        if (channel == null) {
            throw new RuntimeException("None of RabbitMQ channels are available");
        }
        setupQueue();
        consumer = new QueueingConsumer(channel);
        RuntimeContext runtimeContext = getRuntimeContext();
        if (runtimeContext instanceof StreamingRuntimeContext && ((StreamingRuntimeContext) runtimeContext).isCheckpointingEnabled()) {
            autoAck = false;
            // enables transaction mode
            channel.txSelect();
        } else {
            autoAck = true;
        }
        LOG.debug("Starting RabbitMQ source with autoAck status: " + autoAck);
        channel.basicConsume(queueName, autoAck, consumer);
    } catch (IOException e) {
        IOUtils.closeAllQuietly(channel, connection);
        throw new RuntimeException("Cannot create RMQ connection with " + queueName + " at " + rmqConnectionConfig.getHost(), e);
    }
    this.deliveryDeserializer.open(RuntimeContextInitializationContextAdapters.deserializationAdapter(getRuntimeContext(), metricGroup -> metricGroup.addGroup("user")));
    running = true;
}