use of org.apache.flink.connector.testutils.source.reader.TestingReaderContext in project flink by apache.
the class FileSourceReaderTest method testRequestSplitWhenNoSplitRestored.
@Test
public void testRequestSplitWhenNoSplitRestored() throws Exception {
final TestingReaderContext context = new TestingReaderContext();
final FileSourceReader<String, FileSourceSplit> reader = createReader(context);
reader.start();
reader.close();
assertEquals(1, context.getNumSplitRequests());
}
use of org.apache.flink.connector.testutils.source.reader.TestingReaderContext in project flink by apache.
the class FileSourceReaderTest method testNoSplitRequestWhenSplitRestored.
@Test
public void testNoSplitRequestWhenSplitRestored() throws Exception {
final TestingReaderContext context = new TestingReaderContext();
final FileSourceReader<String, FileSourceSplit> reader = createReader(context);
reader.addSplits(Collections.singletonList(createTestFileSplit()));
reader.start();
reader.close();
assertEquals(0, context.getNumSplitRequests());
}
use of org.apache.flink.connector.testutils.source.reader.TestingReaderContext in project flink by apache.
the class HybridSourceReaderTest method testReader.
@Test
public void testReader() throws Exception {
TestingReaderContext readerContext = new TestingReaderContext();
TestingReaderOutput<Integer> readerOutput = new TestingReaderOutput<>();
MockBaseSource source = new MockBaseSource(1, 1, Boundedness.BOUNDED);
// 2 underlying readers to exercise switch
SourceReader<Integer, MockSourceSplit> mockSplitReader1 = source.createReader(readerContext);
SourceReader<Integer, MockSourceSplit> mockSplitReader2 = source.createReader(readerContext);
HybridSourceReader<Integer> reader = new HybridSourceReader<>(readerContext);
Assert.assertThat(readerContext.getSentEvents(), Matchers.emptyIterable());
reader.start();
assertAndClearSourceReaderFinishedEvent(readerContext, -1);
Assert.assertNull(currentReader(reader));
Assert.assertEquals(InputStatus.NOTHING_AVAILABLE, reader.pollNext(readerOutput));
Source source1 = new MockSource(null, 0) {
@Override
public SourceReader<Integer, MockSourceSplit> createReader(SourceReaderContext readerContext) {
return mockSplitReader1;
}
};
reader.handleSourceEvents(new SwitchSourceEvent(0, source1, false));
MockSourceSplit mockSplit = new MockSourceSplit(0, 0, 1);
mockSplit.addRecord(0);
SwitchedSources switchedSources = new SwitchedSources();
switchedSources.put(0, source);
HybridSourceSplit hybridSplit = HybridSourceSplit.wrapSplit(mockSplit, 0, switchedSources);
reader.addSplits(Collections.singletonList(hybridSplit));
// drain splits
InputStatus status = reader.pollNext(readerOutput);
while (readerOutput.getEmittedRecords().isEmpty() || status == InputStatus.MORE_AVAILABLE) {
status = reader.pollNext(readerOutput);
Thread.sleep(10);
}
Assert.assertThat(readerOutput.getEmittedRecords(), Matchers.contains(0));
reader.pollNext(readerOutput);
Assert.assertEquals("before notifyNoMoreSplits", InputStatus.NOTHING_AVAILABLE, reader.pollNext(readerOutput));
reader.notifyNoMoreSplits();
reader.pollNext(readerOutput);
assertAndClearSourceReaderFinishedEvent(readerContext, 0);
Assert.assertEquals("reader before switch source event", mockSplitReader1, currentReader(reader));
Source source2 = new MockSource(null, 0) {
@Override
public SourceReader<Integer, MockSourceSplit> createReader(SourceReaderContext readerContext) {
return mockSplitReader2;
}
};
reader.handleSourceEvents(new SwitchSourceEvent(1, source2, true));
Assert.assertEquals("reader after switch source event", mockSplitReader2, currentReader(reader));
reader.notifyNoMoreSplits();
Assert.assertEquals("reader 1 after notifyNoMoreSplits", InputStatus.END_OF_INPUT, reader.pollNext(readerOutput));
reader.close();
}
use of org.apache.flink.connector.testutils.source.reader.TestingReaderContext in project flink by apache.
the class SourceReaderBaseTest method testMultipleSplitsWithDifferentFinishingMoments.
@Test
void testMultipleSplitsWithDifferentFinishingMoments() throws Exception {
FutureCompletingBlockingQueue<RecordsWithSplitIds<int[]>> elementsQueue = new FutureCompletingBlockingQueue<>();
MockSplitReader mockSplitReader = MockSplitReader.newBuilder().setNumRecordsPerSplitPerFetch(2).setSeparatedFinishedRecord(false).setBlockingFetch(false).build();
MockSourceReader reader = new MockSourceReader(elementsQueue, () -> mockSplitReader, getConfig(), new TestingReaderContext());
reader.start();
List<MockSourceSplit> splits = Arrays.asList(getSplit(0, 10, Boundedness.BOUNDED), getSplit(1, 12, Boundedness.BOUNDED));
reader.addSplits(splits);
reader.notifyNoMoreSplits();
while (true) {
InputStatus status = reader.pollNext(new TestingReaderOutput<>());
if (status == InputStatus.END_OF_INPUT) {
break;
}
if (status == InputStatus.NOTHING_AVAILABLE) {
reader.isAvailable().get();
}
}
}
use of org.apache.flink.connector.testutils.source.reader.TestingReaderContext in project flink by apache.
the class SourceReaderBaseTest method testPollNextReturnMoreAvailableWhenAllSplitFetcherCloseWithLeftoverElementInQueue.
@Test
void testPollNextReturnMoreAvailableWhenAllSplitFetcherCloseWithLeftoverElementInQueue() throws Exception {
FutureCompletingBlockingQueue<RecordsWithSplitIds<int[]>> elementsQueue = new FutureCompletingBlockingQueue<>();
MockSplitReader mockSplitReader = MockSplitReader.newBuilder().setNumRecordsPerSplitPerFetch(1).setBlockingFetch(true).build();
BlockingShutdownSplitFetcherManager<int[], MockSourceSplit> splitFetcherManager = new BlockingShutdownSplitFetcherManager<>(elementsQueue, () -> mockSplitReader);
final MockSourceReader sourceReader = new MockSourceReader(elementsQueue, splitFetcherManager, getConfig(), new TestingReaderContext());
// Create and add a split that only contains one record
final MockSourceSplit split = new MockSourceSplit(0, 0, 1);
sourceReader.addSplits(Collections.singletonList(split));
sourceReader.notifyNoMoreSplits();
// Add the last record to the split when the splitFetcherManager shutting down SplitFetchers
splitFetcherManager.getInShutdownSplitFetcherFuture().thenRun(() -> split.addRecord(1));
assertThat(sourceReader.pollNext(new TestingReaderOutput<>())).isEqualTo(InputStatus.MORE_AVAILABLE);
}
Aggregations