Code review comment: Yes, that's a good addition.

Code before review:
private static void testSourceDoesNotShutdown(boolean shouldHaveReaders) throws Exception {
final int parallelism = 2;
FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class);
TestCountingSource source = new TestCountingSource(20).withoutSplitting();
UnboundedSourceWrapper<KV<Integer, Integer>, TestCountingSource.CounterMark> sourceWrapper =
new UnboundedSourceWrapper<>("noReader", options, source, parallelism);
StreamingRuntimeContext mock = Mockito.mock(StreamingRuntimeContext.class);
if (shouldHaveReaders) {
Mockito.when(mock.getIndexOfThisSubtask()).thenReturn(0);
} else {
Mockito.when(mock.getIndexOfThisSubtask()).thenReturn(parallelism - 1);
}
Mockito.when(mock.getNumberOfParallelSubtasks()).thenReturn(parallelism);
Mockito.when(mock.getExecutionConfig()).thenReturn(new ExecutionConfig());
ProcessingTimeService timerService = Mockito.mock(ProcessingTimeService.class);
Mockito.when(timerService.getCurrentProcessingTime()).thenReturn(Long.MAX_VALUE);
Mockito.when(mock.getProcessingTimeService()).thenReturn(timerService);
sourceWrapper.setRuntimeContext(mock);
sourceWrapper.open(new Configuration());
SourceFunction.SourceContext sourceContext = Mockito.mock(SourceFunction.SourceContext.class);
Object checkpointLock = new Object();
Mockito.when(sourceContext.getCheckpointLock()).thenReturn(checkpointLock);
sourceWrapper.setSourceContext(sourceContext);
sourceWrapper.open(new Configuration());
assertThat(sourceWrapper.getLocalReaders().isEmpty(), is(!shouldHaveReaders));
Thread thread =
new Thread(
() -> {
try {
sourceWrapper.run(sourceContext);
} catch (Exception e) {
LOG.error("Error while running UnboundedSourceWrapper", e);
}
});
try {
thread.start();
if (!shouldHaveReaders) {
while (true) {
StackTraceElement[] callStack = thread.getStackTrace();
if (callStack.length >= 2
&& "sleep".equals(callStack[0].getMethodName())
&& "finalizeSource".equals(callStack[1].getMethodName())) {
break;
}
Thread.sleep(10);
}
}
assertThat(sourceWrapper.isRunning(), is(true));
synchronized (checkpointLock) {
sourceWrapper.onProcessingTime(42);
}
assertThat(sourceWrapper.isRunning(), is(true));
assertThat(thread.isAlive(), is(true));
sourceWrapper.cancel();
} finally {
thread.interrupt();
thread.join();
}
}

Line under review: thread.interrupt();

Code after addressing the review:
private static void testSourceDoesNotShutdown(boolean shouldHaveReaders) throws Exception {
final int parallelism = 2;
FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class);
TestCountingSource source = new TestCountingSource(20).withoutSplitting();
UnboundedSourceWrapper<KV<Integer, Integer>, TestCountingSource.CounterMark> sourceWrapper =
new UnboundedSourceWrapper<>("noReader", options, source, parallelism);
StreamingRuntimeContext mock = Mockito.mock(StreamingRuntimeContext.class);
if (shouldHaveReaders) {
Mockito.when(mock.getIndexOfThisSubtask()).thenReturn(0);
} else {
Mockito.when(mock.getIndexOfThisSubtask()).thenReturn(parallelism - 1);
}
Mockito.when(mock.getNumberOfParallelSubtasks()).thenReturn(parallelism);
Mockito.when(mock.getExecutionConfig()).thenReturn(new ExecutionConfig());
ProcessingTimeService timerService = Mockito.mock(ProcessingTimeService.class);
Mockito.when(timerService.getCurrentProcessingTime()).thenReturn(Long.MAX_VALUE);
Mockito.when(mock.getProcessingTimeService()).thenReturn(timerService);
sourceWrapper.setRuntimeContext(mock);
sourceWrapper.open(new Configuration());
SourceFunction.SourceContext sourceContext = Mockito.mock(SourceFunction.SourceContext.class);
Object checkpointLock = new Object();
Mockito.when(sourceContext.getCheckpointLock()).thenReturn(checkpointLock);
sourceWrapper.setSourceContext(sourceContext);
sourceWrapper.open(new Configuration());
assertThat(sourceWrapper.getLocalReaders().isEmpty(), is(!shouldHaveReaders));
Thread thread =
new Thread(
() -> {
try {
sourceWrapper.run(sourceContext);
} catch (Exception e) {
LOG.error("Error while running UnboundedSourceWrapper", e);
}
});
try {
thread.start();
if (!shouldHaveReaders) {
while (true) {
StackTraceElement[] callStack = thread.getStackTrace();
if (callStack.length >= 2
&& "sleep".equals(callStack[0].getMethodName())
&& "finalizeSource".equals(callStack[1].getMethodName())) {
break;
}
Thread.sleep(10);
}
}
assertThat(sourceWrapper.isRunning(), is(true));
synchronized (checkpointLock) {
sourceWrapper.onProcessingTime(42);
}
assertThat(sourceWrapper.isRunning(), is(true));
assertThat(thread.isAlive(), is(true));
sourceWrapper.cancel();
} finally {
thread.interrupt();
thread.join(1000);
assertThat(thread.isAlive(), is(false));
}
}