org.apache.flink.streaming.runtime.tasks.TestProcessingTimeService Java Examples

private MockFetcher(Map<String, MessageId>... stateSnapshotsToReturn) throws Exception {
    super(
            new TestSourceContext<>(),
            new HashMap<>(),
            null,
            null,
            new TestProcessingTimeService(),
            0,
            MockFetcher.class.getClassLoader(),
            null,
            null,
            null,
            0,
            null,
            null);

    this.stateSnapshotsToReturn.addAll(Arrays.asList(stateSnapshotsToReturn));
}
 

@Test
public void testKeyGroupStartIndexSetting() {

	int startKeyGroupIdx = 7;
	int endKeyGroupIdx = 21;
	KeyGroupRange testKeyGroupList = new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx);

	TestKeyContext keyContext = new TestKeyContext();

	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();

	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		testKeyGroupList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		createQueueFactory());

	Assert.assertEquals(startKeyGroupIdx, service.getLocalKeyGroupRangeStartIdx());
}
 

@Test
public void testCurrentProcessingTime() throws Exception {

	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	processingTimeService.setCurrentTime(17L);
	assertEquals(17, timerService.currentProcessingTime());

	processingTimeService.setCurrentTime(42);
	assertEquals(42, timerService.currentProcessingTime());
}
 

@Test
public void testCurrentEventTime() throws Exception {

	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	timerService.advanceWatermark(17);
	assertEquals(17, timerService.currentWatermark());

	timerService.advanceWatermark(42);
	assertEquals(42, timerService.currentWatermark());
}
 

@Test
public void testCurrentProcessingTime() throws Exception {

	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	processingTimeService.setCurrentTime(17L);
	assertEquals(17, timerService.currentProcessingTime());

	processingTimeService.setCurrentTime(42);
	assertEquals(42, timerService.currentProcessingTime());
}
 

@Test
public void testKeyGroupStartIndexSetting() {

	int startKeyGroupIdx = 7;
	int endKeyGroupIdx = 21;
	KeyGroupRange testKeyGroupList = new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx);

	TestKeyContext keyContext = new TestKeyContext();

	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();

	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		testKeyGroupList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		createQueueFactory());

	Assert.assertEquals(startKeyGroupIdx, service.getLocalKeyGroupRangeStartIdx());
}
 

@Test
public void testKeyGroupStartIndexSetting() {

	int startKeyGroupIdx = 7;
	int endKeyGroupIdx = 21;
	KeyGroupRange testKeyGroupList = new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx);

	TestKeyContext keyContext = new TestKeyContext();

	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();

	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		testKeyGroupList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		createQueueFactory());

	Assert.assertEquals(startKeyGroupIdx, service.getLocalKeyGroupRangeStartIdx());
}
 

@Test
public void testCurrentProcessingTime() throws Exception {

	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	processingTimeService.setCurrentTime(17L);
	assertEquals(17, timerService.currentProcessingTime());

	processingTimeService.setCurrentTime(42);
	assertEquals(42, timerService.currentProcessingTime());
}
 

@Test
public void testIgnorePartitionStateSentinelInSnapshot() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(new KafkaTopicPartition(testTopic, 1), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);
	originalPartitions.put(new KafkaTopicPartition(testTopic, 2), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
	originalPartitions.put(new KafkaTopicPartition(testTopic, 3), KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestFetcher<Long> fetcher = new TestFetcher<>(
			sourceContext,
			originalPartitions,
			null, /* watermark strategy */
			new TestProcessingTimeService(),
			0);

	synchronized (sourceContext.getCheckpointLock()) {
		HashMap<KafkaTopicPartition, Long> currentState = fetcher.snapshotCurrentState();
		fetcher.commitInternalOffsetsToKafka(currentState, new KafkaCommitCallback() {
			@Override
			public void onSuccess() {
			}

			@Override
			public void onException(Throwable cause) {
				throw new RuntimeException("Callback failed", cause);
			}
		});

		assertTrue(fetcher.getLastCommittedOffsets().isPresent());
		assertEquals(Collections.emptyMap(), fetcher.getLastCommittedOffsets().get());
	}
}
 

public TestingSourceOperator(
		SourceReader<T, MockSourceSplit> reader,
		OperatorEventGateway eventGateway,
		int subtaskIndex) {

	this(reader, WatermarkStrategy.noWatermarks(), new TestProcessingTimeService(), eventGateway, subtaskIndex, 5);
}
 

@Test
public void testSkipCorruptedRecordWithPunctuatedWatermarks() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(new KafkaTopicPartition(testTopic, 1), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestProcessingTimeService processingTimeProvider = new TestProcessingTimeService();

	TestFetcher<Long> fetcher = new TestFetcher<>(
		sourceContext,
		originalPartitions,
		null, /* periodic watermark assigner */
		new SerializedValue<AssignerWithPunctuatedWatermarks<Long>>(new PunctuatedTestExtractor()), /* punctuated watermark assigner */
		processingTimeProvider,
		0);

	final KafkaTopicPartitionState<Object> partitionStateHolder = fetcher.subscribedPartitionStates().get(0);

	// elements generate a watermark if the timestamp is a multiple of three
	fetcher.emitRecord(1L, partitionStateHolder, 1L);
	fetcher.emitRecord(2L, partitionStateHolder, 2L);
	fetcher.emitRecord(3L, partitionStateHolder, 3L);
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());
	assertTrue(sourceContext.hasWatermark());
	assertEquals(3L, sourceContext.getLatestWatermark().getTimestamp());
	assertEquals(3L, partitionStateHolder.getOffset());

	// emit null record
	fetcher.emitRecord(null, partitionStateHolder, 4L);

	// no elements or watermarks should have been collected
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());
	assertFalse(sourceContext.hasWatermark());
	// the offset in state still should have advanced
	assertEquals(4L, partitionStateHolder.getOffset());
}
 

@Test
public void testConcurrentPartitionsDiscoveryAndLoopFetching() throws Exception {
	// test data
	final KafkaTopicPartition testPartition = new KafkaTopicPartition("test", 42);

	// ----- create the test fetcher -----

	SourceContext<String> sourceContext = new TestSourceContext<>();
	Map<KafkaTopicPartition, Long> partitionsWithInitialOffsets =
		Collections.singletonMap(testPartition, KafkaTopicPartitionStateSentinel.GROUP_OFFSET);

	final OneShotLatch fetchLoopWaitLatch = new OneShotLatch();
	final OneShotLatch stateIterationBlockLatch = new OneShotLatch();

	final TestFetcher<String> fetcher = new TestFetcher<>(
		sourceContext,
		partitionsWithInitialOffsets,
		null, /* watermark strategy */
		new TestProcessingTimeService(),
		10,
		fetchLoopWaitLatch,
		stateIterationBlockLatch);

	// ----- run the fetcher -----

	final CheckedThread checkedThread = new CheckedThread() {
		@Override
		public void go() throws Exception {
			fetcher.runFetchLoop();
		}
	};
	checkedThread.start();

	// wait until state iteration begins before adding discovered partitions
	fetchLoopWaitLatch.await();
	fetcher.addDiscoveredPartitions(Collections.singletonList(testPartition));

	stateIterationBlockLatch.trigger();
	checkedThread.sync();
}
 

@Test
public void testSkipCorruptedRecordWithPunctuatedWatermarks() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(new KafkaTopicPartition(testTopic, 1), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestProcessingTimeService processingTimeProvider = new TestProcessingTimeService();

	TestFetcher<Long> fetcher = new TestFetcher<>(
		sourceContext,
		originalPartitions,
		null, /* periodic watermark assigner */
		new SerializedValue<AssignerWithPunctuatedWatermarks<Long>>(new PunctuatedTestExtractor()), /* punctuated watermark assigner */
		processingTimeProvider,
		0);

	final KafkaTopicPartitionState<Object> partitionStateHolder = fetcher.subscribedPartitionStates().get(0);

	// elements generate a watermark if the timestamp is a multiple of three
	fetcher.emitRecord(1L, partitionStateHolder, 1L);
	fetcher.emitRecord(2L, partitionStateHolder, 2L);
	fetcher.emitRecord(3L, partitionStateHolder, 3L);
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());
	assertTrue(sourceContext.hasWatermark());
	assertEquals(3L, sourceContext.getLatestWatermark().getTimestamp());
	assertEquals(3L, partitionStateHolder.getOffset());

	// emit null record
	fetcher.emitRecord(null, partitionStateHolder, 4L);

	// no elements or watermarks should have been collected
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());
	assertFalse(sourceContext.hasWatermark());
	// the offset in state still should have advanced
	assertEquals(4L, partitionStateHolder.getOffset());
}
 

@SafeVarargs
private MockFetcher(HashMap<KafkaTopicPartition, Long>... stateSnapshotsToReturn) throws Exception {
	super(
			new TestSourceContext<>(),
			new HashMap<>(),
			null,
			null,
			new TestProcessingTimeService(),
			0,
			MockFetcher.class.getClassLoader(),
			new UnregisteredMetricsGroup(),
			false);

	this.stateSnapshotsToReturn.addAll(Arrays.asList(stateSnapshotsToReturn));
}
 

@Test
public void testPeriodicWatermarksWithNoSubscribedPartitionsShouldYieldNoWatermarks() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestProcessingTimeService processingTimeProvider = new TestProcessingTimeService();

	TestFetcher<Long> fetcher = new TestFetcher<>(
			sourceContext,
			originalPartitions,
			new SerializedValue<>(testWmStrategy),
			processingTimeProvider,
			10);

	processingTimeProvider.setCurrentTime(10);
	// no partitions; when the periodic watermark emitter fires, no watermark should be emitted
	assertFalse(sourceContext.hasWatermark());

	// counter-test that when the fetcher does actually have partitions,
	// when the periodic watermark emitter fires again, a watermark really is emitted
	fetcher.addDiscoveredPartitions(Collections.singletonList(
			new KafkaTopicPartition(testTopic, 0)));
	emitRecord(fetcher, 100L, fetcher.subscribedPartitionStates().get(0), 3L);
	processingTimeProvider.setCurrentTime(20);
	assertEquals(100, sourceContext.getLatestWatermark().getTimestamp());
}
 

@SafeVarargs
private MockFetcher(HashMap<KafkaTopicPartition, Long>... stateSnapshotsToReturn) throws Exception {
	super(
			new TestSourceContext<>(),
			new HashMap<>(),
			null,
			null,
			new TestProcessingTimeService(),
			0,
			MockFetcher.class.getClassLoader(),
			new UnregisteredMetricsGroup(),
			false);

	this.stateSnapshotsToReturn.addAll(Arrays.asList(stateSnapshotsToReturn));
}
 

@Test
public void testSkipCorruptedRecord() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(new KafkaTopicPartition(testTopic, 1), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestFetcher<Long> fetcher = new TestFetcher<>(
			sourceContext,
			originalPartitions,
			null, /* watermark strategy */
			new TestProcessingTimeService(),
			0);

	final KafkaTopicPartitionState<Long, Object> partitionStateHolder = fetcher.subscribedPartitionStates().get(0);

	emitRecord(fetcher, 1L, partitionStateHolder, 1L);
	emitRecord(fetcher, 2L, partitionStateHolder, 2L);
	assertEquals(2L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(2L, partitionStateHolder.getOffset());

	// emit no records
	fetcher.emitRecordsWithTimestamps(emptyQueue(), partitionStateHolder, 3L, Long.MIN_VALUE);
	assertEquals(2L, sourceContext.getLatestElement().getValue().longValue()); // the null record should be skipped
	assertEquals(3L, partitionStateHolder.getOffset()); // the offset in state still should have advanced
}
 

/**
 * This also verifies that we iterate over all timers and set the key context on each element.
 */
@Test
public void testForEachProcessingTimeTimers() throws Exception {
	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
		createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	// get two different keys
	int key1 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	int key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	while (key2 == key1) {
		key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	}

	Set<Tuple3<Integer, String, Long>> timers = new HashSet<>();
	timers.add(Tuple3.of(key1, "ciao", 10L));
	timers.add(Tuple3.of(key1, "hello", 10L));
	timers.add(Tuple3.of(key2, "ciao", 10L));
	timers.add(Tuple3.of(key2, "hello", 10L));

	for (Tuple3<Integer, String, Long> timer : timers) {
		keyContext.setCurrentKey(timer.f0);
		timerService.registerProcessingTimeTimer(timer.f1, timer.f2);
	}

	Set<Tuple3<Integer, String, Long>> results = new HashSet<>();
	timerService.forEachProcessingTimeTimer((namespace, timer) -> {
		results.add(Tuple3.of((Integer) keyContext.getCurrentKey(), namespace, timer));
	});

	Assert.assertEquals(timers, results);
}
 

@Test
public void testAutomaticWatermarkContext() throws Exception {

	// regular stream source operator
	final StreamSource<String, InfiniteSource<String>> operator =
		new StreamSource<>(new InfiniteSource<>());

	long watermarkInterval = 10;
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	processingTimeService.setCurrentTime(0);

	setupSourceOperator(operator, TimeCharacteristic.IngestionTime, watermarkInterval, processingTimeService);

	final List<StreamElement> output = new ArrayList<>();

	StreamSourceContexts.getSourceContext(TimeCharacteristic.IngestionTime,
		operator.getContainingTask().getProcessingTimeService(),
		operator.getContainingTask().getCheckpointLock(),
		operator.getContainingTask().getStreamStatusMaintainer(),
		new CollectorOutput<String>(output),
		operator.getExecutionConfig().getAutoWatermarkInterval(),
		-1);

	// periodically emit the watermarks
	// even though we start from 1 the watermark are still
	// going to be aligned with the watermark interval.

	for (long i = 1; i < 100; i += watermarkInterval)  {
		processingTimeService.setCurrentTime(i);
	}

	assertTrue(output.size() == 9);

	long nextWatermark = 0;
	for (StreamElement el : output) {
		nextWatermark += watermarkInterval;
		Watermark wm = (Watermark) el;
		assertTrue(wm.getTimestamp() == nextWatermark);
	}
}
 

@Test
public void testIgnorePartitionStateSentinelInSnapshot() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(new KafkaTopicPartition(testTopic, 1), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);
	originalPartitions.put(new KafkaTopicPartition(testTopic, 2), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
	originalPartitions.put(new KafkaTopicPartition(testTopic, 3), KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestFetcher<Long> fetcher = new TestFetcher<>(
		sourceContext,
		originalPartitions,
		null,
		null,
		new TestProcessingTimeService(),
		0);

	synchronized (sourceContext.getCheckpointLock()) {
		HashMap<KafkaTopicPartition, Long> currentState = fetcher.snapshotCurrentState();
		fetcher.commitInternalOffsetsToKafka(currentState, new KafkaCommitCallback() {
			@Override
			public void onSuccess() {
			}

			@Override
			public void onException(Throwable cause) {
				throw new RuntimeException("Callback failed", cause);
			}
		});

		assertTrue(fetcher.getLastCommittedOffsets().isPresent());
		assertEquals(Collections.emptyMap(), fetcher.getLastCommittedOffsets().get());
	}
}
 

@SafeVarargs
private MockFetcher(HashMap<KafkaTopicPartition, Long>... stateSnapshotsToReturn) throws Exception {
	super(
			new TestSourceContext<>(),
			new HashMap<>(),
			null /* watermark strategy */,
			new TestProcessingTimeService(),
			0,
			MockFetcher.class.getClassLoader(),
			new UnregisteredMetricsGroup(),
			false);

	this.stateSnapshotsToReturn.addAll(Arrays.asList(stateSnapshotsToReturn));
}
 

@Test
public void testIgnorePartitionStateSentinelInSnapshot() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(new KafkaTopicPartition(testTopic, 1), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);
	originalPartitions.put(new KafkaTopicPartition(testTopic, 2), KafkaTopicPartitionStateSentinel.GROUP_OFFSET);
	originalPartitions.put(new KafkaTopicPartition(testTopic, 3), KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestFetcher<Long> fetcher = new TestFetcher<>(
		sourceContext,
		originalPartitions,
		null,
		null,
		new TestProcessingTimeService(),
		0);

	synchronized (sourceContext.getCheckpointLock()) {
		HashMap<KafkaTopicPartition, Long> currentState = fetcher.snapshotCurrentState();
		fetcher.commitInternalOffsetsToKafka(currentState, new KafkaCommitCallback() {
			@Override
			public void onSuccess() {
			}

			@Override
			public void onException(Throwable cause) {
				throw new RuntimeException("Callback failed", cause);
			}
		});

		assertTrue(fetcher.getLastCommittedOffsets().isPresent());
		assertEquals(Collections.emptyMap(), fetcher.getLastCommittedOffsets().get());
	}
}
 

@Test
public void testConcurrentPartitionsDiscoveryAndLoopFetching() throws Exception {
    String tp = topicName("test", 2);
    TestSourceContext<Long> sourceContext = new TestSourceContext<Long>();
    Map<String, MessageId> offset = Collections.singletonMap(topicName(tp, 1), MessageId.latest);

    OneShotLatch fetchLoopWaitLatch = new OneShotLatch();
    OneShotLatch stateIterationBlockLatch = new OneShotLatch();

    TestFetcher fetcher = new TestFetcher(
            sourceContext,
            offset,
            null,
            null,
            new TestProcessingTimeService(),
            10,
            fetchLoopWaitLatch,
            stateIterationBlockLatch);

    // ----- run the fetcher -----

    final CheckedThread checkedThread = new CheckedThread() {
        @Override
        public void go() throws Exception {
            fetcher.runFetchLoop();
        }
    };
    checkedThread.start();

    fetchLoopWaitLatch.await();
    fetcher.addDiscoveredTopics(Sets.newSet(tp));

    stateIterationBlockLatch.trigger();
    checkedThread.sync();
}
 

@Test
public void testSkipCorruptedRecord() throws Exception {
    String testTopic = "tp";
    Map<String, MessageId> offset = Collections.singletonMap(topicName(testTopic, 1), MessageId.latest);

    TestSourceContext<Long> sourceContext = new TestSourceContext<Long>();
    TestFetcher<Long> fetcher = new TestFetcher<>(
            sourceContext,
            offset,
            null,
            null,
            new TestProcessingTimeService(),
            0,
            null,
            null);

    PulsarTopicState stateHolder = fetcher.getSubscribedTopicStates().get(0);
    fetcher.emitRecord(1L, stateHolder, dummyMessageId(1));
    fetcher.emitRecord(2L, stateHolder, dummyMessageId(2));
    assertEquals(2L, sourceContext.getLatestElement().getValue().longValue());
    assertEquals(dummyMessageId(2), stateHolder.getOffset());

    // emit null record
    fetcher.emitRecord(null, stateHolder, dummyMessageId(3));
    assertEquals(2L, sourceContext.getLatestElement().getValue().longValue());
    assertEquals(dummyMessageId(3), stateHolder.getOffset());
}
 

@Test
public void testIgnorePartitionStates() throws Exception {
    String testTopic = "tp";
    Map<String, MessageId> offset = new HashMap<>();
    offset.put(topicName(testTopic, 1), MessageId.earliest);
    offset.put(topicName(testTopic, 2), MessageId.latest);

    TestSourceContext<Long> sourceContext = new TestSourceContext<>();
    TestFetcher<Long> fetcher = new TestFetcher<>(
            sourceContext,
            offset,
            null,
            null,
            new TestProcessingTimeService(),
            0,
            null,
            null);

    synchronized (sourceContext.getCheckpointLock()) {
        Map<String, MessageId> current = fetcher.snapshotCurrentState();
        fetcher.commitOffsetToPulsar(current, new PulsarCommitCallback() {
            @Override
            public void onSuccess() {
            }

            @Override
            public void onException(Throwable cause) {
                throw new RuntimeException("Callback failed", cause);
            }
        });

        assertTrue(fetcher.lastCommittedOffsets.isPresent());
        assertEquals(fetcher.lastCommittedOffsets.get().size(), 0);
    }
}
 

/**
 * This also verifies that we iterate over all timers and set the key context on each element.
 */
@Test
public void testForEachProcessingTimeTimers() throws Exception {
	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
		createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	// get two different keys
	int key1 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	int key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	while (key2 == key1) {
		key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	}

	Set<Tuple3<Integer, String, Long>> timers = new HashSet<>();
	timers.add(Tuple3.of(key1, "ciao", 10L));
	timers.add(Tuple3.of(key1, "hello", 10L));
	timers.add(Tuple3.of(key2, "ciao", 10L));
	timers.add(Tuple3.of(key2, "hello", 10L));

	for (Tuple3<Integer, String, Long> timer : timers) {
		keyContext.setCurrentKey(timer.f0);
		timerService.registerProcessingTimeTimer(timer.f1, timer.f2);
	}

	Set<Tuple3<Integer, String, Long>> results = new HashSet<>();
	timerService.forEachProcessingTimeTimer((namespace, timer) -> {
		results.add(Tuple3.of((Integer) keyContext.getCurrentKey(), namespace, timer));
	});

	Assert.assertEquals(timers, results);
}
 

/**
 * This also verifies that we iterate over all timers and set the key context on each element.
 */
@Test
public void testForEachEventTimeTimers() throws Exception {
	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
		createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	// get two different keys
	int key1 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	int key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	while (key2 == key1) {
		key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	}

	Set<Tuple3<Integer, String, Long>> timers = new HashSet<>();
	timers.add(Tuple3.of(key1, "ciao", 10L));
	timers.add(Tuple3.of(key1, "hello", 10L));
	timers.add(Tuple3.of(key2, "ciao", 10L));
	timers.add(Tuple3.of(key2, "hello", 10L));

	for (Tuple3<Integer, String, Long> timer : timers) {
		keyContext.setCurrentKey(timer.f0);
		timerService.registerEventTimeTimer(timer.f1, timer.f2);
	}

	Set<Tuple3<Integer, String, Long>> results = new HashSet<>();
	timerService.forEachEventTimeTimer((namespace, timer) -> {
		results.add(Tuple3.of((Integer) keyContext.getCurrentKey(), namespace, timer));
	});

	Assert.assertEquals(timers, results);
}
 

@Test
public void testPeriodicWatermarksWithNoSubscribedPartitionsShouldYieldNoWatermarks() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestProcessingTimeService processingTimeProvider = new TestProcessingTimeService();

	TestFetcher<Long> fetcher = new TestFetcher<>(
		sourceContext,
		originalPartitions,
		new SerializedValue<AssignerWithPeriodicWatermarks<Long>>(new PeriodicTestExtractor()),
		null, /* punctuated watermarks assigner*/
		processingTimeProvider,
		10);

	processingTimeProvider.setCurrentTime(10);
	// no partitions; when the periodic watermark emitter fires, no watermark should be emitted
	assertFalse(sourceContext.hasWatermark());

	// counter-test that when the fetcher does actually have partitions,
	// when the periodic watermark emitter fires again, a watermark really is emitted
	fetcher.addDiscoveredPartitions(Collections.singletonList(new KafkaTopicPartition(testTopic, 0)));
	fetcher.emitRecord(100L, fetcher.subscribedPartitionStates().get(0), 3L);
	processingTimeProvider.setCurrentTime(20);
	assertEquals(100, sourceContext.getLatestWatermark().getTimestamp());
}
 

/**
 * This also verifies that we iterate over all timers and set the key context on each element.
 */
@Test
public void testForEachEventTimeTimers() throws Exception {
	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
		createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	// get two different keys
	int key1 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	int key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	while (key2 == key1) {
		key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	}

	Set<Tuple3<Integer, String, Long>> timers = new HashSet<>();
	timers.add(Tuple3.of(key1, "ciao", 10L));
	timers.add(Tuple3.of(key1, "hello", 10L));
	timers.add(Tuple3.of(key2, "ciao", 10L));
	timers.add(Tuple3.of(key2, "hello", 10L));

	for (Tuple3<Integer, String, Long> timer : timers) {
		keyContext.setCurrentKey(timer.f0);
		timerService.registerEventTimeTimer(timer.f1, timer.f2);
	}

	Set<Tuple3<Integer, String, Long>> results = new HashSet<>();
	timerService.forEachEventTimeTimer((namespace, timer) -> {
		results.add(Tuple3.of((Integer) keyContext.getCurrentKey(), namespace, timer));
	});

	Assert.assertEquals(timers, results);
}
 

void setupSourceOperator(
	StreamSource<Long, ?> operator,
	TestProcessingTimeService testProcessingTimeService
);