Examples with MockSourceSplit org.apache.flink.api.connector.source.mocks.MockSourceSplit used on opensource projects

Example 1 with MockSourceSplit

use of org.apache.flink.api.connector.source.mocks.MockSourceSplit in project flink by apache.

the class HybridSourceSplitEnumeratorTest method testRestoreEnumerator.

@Test
public void testRestoreEnumerator() throws Exception {
    setupEnumeratorAndTriggerSourceSwitch();
    enumerator = (HybridSourceSplitEnumerator) source.createEnumerator(context);
    enumerator.start();
    HybridSourceEnumeratorState enumeratorState = enumerator.snapshotState(0);
    MockSplitEnumerator underlyingEnumerator = getCurrentEnumerator(enumerator);
    Assert.assertThat((List<MockSourceSplit>) Whitebox.getInternalState(underlyingEnumerator, "splits"), Matchers.iterableWithSize(1));
    enumerator = (HybridSourceSplitEnumerator) source.restoreEnumerator(context, enumeratorState);
    enumerator.start();
    underlyingEnumerator = getCurrentEnumerator(enumerator);
    Assert.assertThat((List<MockSourceSplit>) Whitebox.getInternalState(underlyingEnumerator, "splits"), Matchers.iterableWithSize(1));
}

Example 2 with MockSourceSplit

use of org.apache.flink.api.connector.source.mocks.MockSourceSplit 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();
}

Example 3 with MockSourceSplit

use of org.apache.flink.api.connector.source.mocks.MockSourceSplit 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();
        }
    }
}

Example 4 with MockSourceSplit

use of org.apache.flink.api.connector.source.mocks.MockSourceSplit 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);
}

Example 5 with MockSourceSplit

use of org.apache.flink.api.connector.source.mocks.MockSourceSplit in project flink by apache.

the class SourceReaderBaseTest method testExceptionInSplitReader.

@Test
void testExceptionInSplitReader() {
    assertThatThrownBy(() -> {
        final String errMsg = "Testing Exception";
        FutureCompletingBlockingQueue<RecordsWithSplitIds<int[]>> elementsQueue = new FutureCompletingBlockingQueue<>();
        // called.
        try (MockSourceReader reader = new MockSourceReader(elementsQueue, () -> new SplitReader<int[], MockSourceSplit>() {

            @Override
            public RecordsWithSplitIds<int[]> fetch() {
                throw new RuntimeException(errMsg);
            }

            @Override
            public void handleSplitsChanges(SplitsChange<MockSourceSplit> splitsChanges) {
            }

            @Override
            public void wakeUp() {
            }

            @Override
            public void close() {
            }
        }, getConfig(), new TestingReaderContext())) {
            ValidatingSourceOutput output = new ValidatingSourceOutput();
            reader.addSplits(Collections.singletonList(getSplit(0, NUM_RECORDS_PER_SPLIT, Boundedness.CONTINUOUS_UNBOUNDED)));
            reader.notifyNoMoreSplits();
            // two polls.
            while (true) {
                InputStatus inputStatus = reader.pollNext(output);
                assertThat(inputStatus).isNotEqualTo(InputStatus.END_OF_INPUT);
                // Add a sleep to avoid tight loop.
                Thread.sleep(1);
            }
        }
    }).isInstanceOf(RuntimeException.class).hasMessage("One or more fetchers have encountered exception");
}