Examples with MockBaseSource org.apache.flink.connector.base.source.reader.mocks.MockBaseSource used on opensource projects

Example 6 with MockBaseSource

use of org.apache.flink.connector.base.source.reader.mocks.MockBaseSource 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 7 with MockBaseSource

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

the class HybridSourceReaderTest method testReaderRecovery.

@Test
public void testReaderRecovery() throws Exception {
    TestingReaderContext readerContext = new TestingReaderContext();
    TestingReaderOutput<Integer> readerOutput = new TestingReaderOutput<>();
    MockBaseSource source = new MockBaseSource(1, 1, Boundedness.BOUNDED);
    HybridSourceReader<Integer> reader = new HybridSourceReader<>(readerContext);
    reader.start();
    assertAndClearSourceReaderFinishedEvent(readerContext, -1);
    reader.handleSourceEvents(new SwitchSourceEvent(0, source, false));
    MockSourceSplit mockSplit = new MockSourceSplit(0, 0, 2147483647);
    SwitchedSources switchedSources = new SwitchedSources();
    switchedSources.put(0, source);
    HybridSourceSplit hybridSplit = HybridSourceSplit.wrapSplit(mockSplit, 0, switchedSources);
    reader.addSplits(Collections.singletonList(hybridSplit));
    List<HybridSourceSplit> snapshot = reader.snapshotState(0);
    Assert.assertThat(snapshot, Matchers.contains(hybridSplit));
    // reader recovery
    readerContext.clearSentEvents();
    reader = new HybridSourceReader<>(readerContext);
    reader.addSplits(snapshot);
    Assert.assertNull(currentReader(reader));
    reader.start();
    Assert.assertNull(currentReader(reader));
    assertAndClearSourceReaderFinishedEvent(readerContext, -1);
    reader.handleSourceEvents(new SwitchSourceEvent(0, source, false));
    Assert.assertNotNull(currentReader(reader));
    Assert.assertThat(reader.snapshotState(1), Matchers.contains(hybridSplit));
    reader.close();
}

Example 8 with MockBaseSource

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

the class HybridSourceReaderTest method testDefaultMethodDelegation.

@Test
public void testDefaultMethodDelegation() throws Exception {
    TestingReaderContext readerContext = new TestingReaderContext();
    TestingReaderOutput<Integer> readerOutput = new TestingReaderOutput<>();
    MockBaseSource source = new MockBaseSource(1, 1, Boundedness.BOUNDED) {

        @Override
        public SourceReader<Integer, MockSourceSplit> createReader(SourceReaderContext readerContext) {
            return Mockito.spy(super.createReader(readerContext));
        }
    };
    HybridSourceReader<Integer> reader = new HybridSourceReader<>(readerContext);
    reader.start();
    assertAndClearSourceReaderFinishedEvent(readerContext, -1);
    reader.handleSourceEvents(new SwitchSourceEvent(0, source, false));
    SourceReader<Integer, MockSourceSplit> underlyingReader = currentReader(reader);
    reader.notifyCheckpointComplete(1);
    Mockito.verify(underlyingReader).notifyCheckpointComplete(1);
    reader.notifyCheckpointAborted(1);
    Mockito.verify(underlyingReader).notifyCheckpointAborted(1);
    reader.close();
}

Example 9 with MockBaseSource

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

the class HybridSourceTest method testBuilderWithSourceFactory.

@Test
public void testBuilderWithSourceFactory() {
    HybridSource.SourceFactory<Integer, Source<Integer, ?, ?>, MockSplitEnumerator> sourceFactory = new HybridSource.SourceFactory<Integer, Source<Integer, ?, ?>, MockSplitEnumerator>() {

        @Override
        public Source<Integer, ?, ?> create(HybridSource.SourceSwitchContext<MockSplitEnumerator> context) {
            MockSplitEnumerator enumerator = context.getPreviousEnumerator();
            return new MockBaseSource(1, 1, Boundedness.BOUNDED);
        }
    };
    HybridSource<Integer> source = new HybridSource.HybridSourceBuilder<Integer, MockSplitEnumerator>().<MockSplitEnumerator, Source<Integer, ?, ?>>addSource(new MockBaseSource(1, 1, Boundedness.BOUNDED)).addSource(sourceFactory, Boundedness.BOUNDED).build();
    assertNotNull(source);
}

Example 10 with MockBaseSource

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

the class CoordinatedSourceITCase method testEnumeratorReaderCommunication.

@Test
public void testEnumeratorReaderCommunication() throws Exception {
    StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
    MockBaseSource source = new MockBaseSource(2, 10, Boundedness.BOUNDED);
    DataStream<Integer> stream = env.fromSource(source, WatermarkStrategy.noWatermarks(), "TestingSource");
    executeAndVerify(env, stream, 20);
}