Examples with SourceContext org.apache.flink.streaming.api.functions.source.SourceFunction.SourceContext used on opensource projects

Example 1 with SourceContext

use of org.apache.flink.streaming.api.functions.source.SourceFunction.SourceContext in project flink by apache.

the class SpoutCollectorTest method testSpoutStormCollector.

@SuppressWarnings({ "rawtypes", "unchecked" })
@Test
public void testSpoutStormCollector() throws InstantiationException, IllegalAccessException {
    for (int numberOfAttributes = -1; numberOfAttributes < 26; ++numberOfAttributes) {
        final SourceContext flinkCollector = mock(SourceContext.class);
        Tuple flinkTuple = null;
        final Values tuple = new Values();
        SpoutCollector<?> collector;
        final String streamId = "streamId";
        HashMap<String, Integer> attributes = new HashMap<String, Integer>();
        attributes.put(streamId, numberOfAttributes);
        if (numberOfAttributes == -1) {
            collector = new SpoutCollector(attributes, -1, flinkCollector);
            tuple.add(new Integer(this.r.nextInt()));
        } else {
            collector = new SpoutCollector(attributes, -1, flinkCollector);
            flinkTuple = Tuple.getTupleClass(numberOfAttributes).newInstance();
            for (int i = 0; i < numberOfAttributes; ++i) {
                tuple.add(new Integer(this.r.nextInt()));
                flinkTuple.setField(tuple.get(i), i);
            }
        }
        final List<Integer> taskIds;
        final Object messageId = new Integer(this.r.nextInt());
        taskIds = collector.emit(streamId, tuple, messageId);
        Assert.assertNull(taskIds);
        if (numberOfAttributes == -1) {
            verify(flinkCollector).collect(tuple.get(0));
        } else {
            verify(flinkCollector).collect(flinkTuple);
        }
    }
}

Example 2 with SourceContext

use of org.apache.flink.streaming.api.functions.source.SourceFunction.SourceContext in project flink by apache.

the class SpoutCollectorTest method testSpoutStormCollectorWithTaskId.

@SuppressWarnings({ "rawtypes", "unchecked" })
@Test
public void testSpoutStormCollectorWithTaskId() throws InstantiationException, IllegalAccessException {
    for (int numberOfAttributes = 0; numberOfAttributes < 25; ++numberOfAttributes) {
        final SourceContext flinkCollector = mock(SourceContext.class);
        final int taskId = 42;
        final String streamId = "streamId";
        HashMap<String, Integer> attributes = new HashMap<String, Integer>();
        attributes.put(streamId, numberOfAttributes);
        SpoutCollector<?> collector = new SpoutCollector(attributes, taskId, flinkCollector);
        final Values tuple = new Values();
        final Tuple flinkTuple = Tuple.getTupleClass(numberOfAttributes + 1).newInstance();
        for (int i = 0; i < numberOfAttributes; ++i) {
            tuple.add(new Integer(this.r.nextInt()));
            flinkTuple.setField(tuple.get(i), i);
        }
        flinkTuple.setField(taskId, numberOfAttributes);
        final List<Integer> taskIds;
        final Object messageId = new Integer(this.r.nextInt());
        taskIds = collector.emit(streamId, tuple, messageId);
        Assert.assertNull(taskIds);
        verify(flinkCollector).collect(flinkTuple);
    }
}

Example 3 with SourceContext

use of org.apache.flink.streaming.api.functions.source.SourceFunction.SourceContext in project rocketmq-externals by apache.

the class RocketMQSourceTest method testSource.

@Test
public void testSource() throws Exception {
    List<MessageExt> msgFoundList = new ArrayList<>();
    MessageExt messageExt = new MessageExt();
    messageExt.setKeys("keys");
    messageExt.setBody("body data".getBytes());
    messageExt.setBornTimestamp(System.currentTimeMillis());
    msgFoundList.add(messageExt);
    PullResult pullResult = new PullResult(PullStatus.FOUND, 3, 1, 5, msgFoundList);
    when(consumer.fetchConsumeOffset(any(MessageQueue.class), anyBoolean())).thenReturn(2L);
    when(consumer.pull(any(MessageQueue.class), anyString(), anyLong(), anyInt())).thenReturn(pullResult);
    SourceContext context = mock(SourceContext.class);
    when(context.getCheckpointLock()).thenReturn(new Object());
    rocketMQSource.run(context);
    // schedule the pull task
    Set<MessageQueue> set = new HashSet();
    set.add(new MessageQueue(topic, "brk", 1));
    pullConsumerScheduleService.putTask(topic, set);
    MessageExt msg = pullResult.getMsgFoundList().get(0);
    // atLeastOnce: re-pulling immediately when messages found before
    verify(context, atLeastOnce()).collectWithTimestamp(deserializationSchema.deserializeKeyAndValue(msg.getKeys().getBytes(), msg.getBody()), msg.getBornTimestamp());
}

Example 4 with SourceContext

use of org.apache.flink.streaming.api.functions.source.SourceFunction.SourceContext in project pinpoint by naver.

the class AgentStatHandler method handleSimple.

@Override
public void handleSimple(ServerRequest<TBase<?, ?>> serverRequest) {
    final TBase<?, ?> tBase = serverRequest.getData();
    final Map<String, String> metaInfo = new HashMap<>(serverRequest.getHeaderEntity().getEntityAll());
    final RawData rawData = new RawData(tBase, metaInfo);
    final SourceContext sourceContext = roundRobinSourceContext();
    if (sourceContext == null) {
        logger.warn("sourceContext is null.");
        return;
    }
    sourceContext.collect(rawData);
}