Search in sources :

Example 1 with MockSplitReader

use of org.apache.flink.connector.base.source.reader.mocks.MockSplitReader 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();
        }
    }
}
Also used : MockSplitReader(org.apache.flink.connector.base.source.reader.mocks.MockSplitReader) FutureCompletingBlockingQueue(org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue) MockSourceReader(org.apache.flink.connector.base.source.reader.mocks.MockSourceReader) InputStatus(org.apache.flink.core.io.InputStatus) TestingReaderContext(org.apache.flink.connector.testutils.source.reader.TestingReaderContext) TestingRecordsWithSplitIds(org.apache.flink.connector.base.source.reader.mocks.TestingRecordsWithSplitIds) MockSourceSplit(org.apache.flink.api.connector.source.mocks.MockSourceSplit) Test(org.junit.jupiter.api.Test)

Example 2 with MockSplitReader

use of org.apache.flink.connector.base.source.reader.mocks.MockSplitReader 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);
}
Also used : MockSplitReader(org.apache.flink.connector.base.source.reader.mocks.MockSplitReader) FutureCompletingBlockingQueue(org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue) MockSourceReader(org.apache.flink.connector.base.source.reader.mocks.MockSourceReader) TestingReaderContext(org.apache.flink.connector.testutils.source.reader.TestingReaderContext) TestingRecordsWithSplitIds(org.apache.flink.connector.base.source.reader.mocks.TestingRecordsWithSplitIds) MockSourceSplit(org.apache.flink.api.connector.source.mocks.MockSourceSplit) Test(org.junit.jupiter.api.Test)

Example 3 with MockSplitReader

use of org.apache.flink.connector.base.source.reader.mocks.MockSplitReader in project flink by apache.

the class SourceReaderBaseTest method createReader.

// ---------------- helper methods -----------------
@Override
protected MockSourceReader createReader() {
    FutureCompletingBlockingQueue<RecordsWithSplitIds<int[]>> elementsQueue = new FutureCompletingBlockingQueue<>();
    MockSplitReader mockSplitReader = MockSplitReader.newBuilder().setNumRecordsPerSplitPerFetch(2).setBlockingFetch(true).build();
    return new MockSourceReader(elementsQueue, () -> mockSplitReader, getConfig(), new TestingReaderContext());
}
Also used : MockSplitReader(org.apache.flink.connector.base.source.reader.mocks.MockSplitReader) FutureCompletingBlockingQueue(org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue) MockSourceReader(org.apache.flink.connector.base.source.reader.mocks.MockSourceReader) TestingReaderContext(org.apache.flink.connector.testutils.source.reader.TestingReaderContext) TestingRecordsWithSplitIds(org.apache.flink.connector.base.source.reader.mocks.TestingRecordsWithSplitIds)

Example 4 with MockSplitReader

use of org.apache.flink.connector.base.source.reader.mocks.MockSplitReader in project flink by apache.

the class SourceReaderBaseTest method testMultipleSplitsWithSeparatedFinishedRecord.

@Test
void testMultipleSplitsWithSeparatedFinishedRecord() throws Exception {
    FutureCompletingBlockingQueue<RecordsWithSplitIds<int[]>> elementsQueue = new FutureCompletingBlockingQueue<>();
    MockSplitReader mockSplitReader = MockSplitReader.newBuilder().setNumRecordsPerSplitPerFetch(2).setSeparatedFinishedRecord(true).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, 10, 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();
        }
    }
}
Also used : MockSplitReader(org.apache.flink.connector.base.source.reader.mocks.MockSplitReader) FutureCompletingBlockingQueue(org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue) MockSourceReader(org.apache.flink.connector.base.source.reader.mocks.MockSourceReader) InputStatus(org.apache.flink.core.io.InputStatus) TestingReaderContext(org.apache.flink.connector.testutils.source.reader.TestingReaderContext) TestingRecordsWithSplitIds(org.apache.flink.connector.base.source.reader.mocks.TestingRecordsWithSplitIds) MockSourceSplit(org.apache.flink.api.connector.source.mocks.MockSourceSplit) Test(org.junit.jupiter.api.Test)

Aggregations

MockSourceReader (org.apache.flink.connector.base.source.reader.mocks.MockSourceReader)4 MockSplitReader (org.apache.flink.connector.base.source.reader.mocks.MockSplitReader)4 TestingRecordsWithSplitIds (org.apache.flink.connector.base.source.reader.mocks.TestingRecordsWithSplitIds)4 FutureCompletingBlockingQueue (org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue)4 TestingReaderContext (org.apache.flink.connector.testutils.source.reader.TestingReaderContext)4 MockSourceSplit (org.apache.flink.api.connector.source.mocks.MockSourceSplit)3 Test (org.junit.jupiter.api.Test)3 InputStatus (org.apache.flink.core.io.InputStatus)2