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

Example 1 with MockSource

use of org.apache.flink.api.connector.source.mocks.MockSource 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 2 with MockSource

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

the class MultipleInputStreamTaskTest method testWatermarkMetrics.

@Test
@SuppressWarnings("unchecked")
public void testWatermarkMetrics() throws Exception {
    OperatorID mainOperatorId = new OperatorID();
    OperatorID chainedOperatorId = new OperatorID();
    InterceptingOperatorMetricGroup mainOperatorMetricGroup = new InterceptingOperatorMetricGroup();
    InterceptingOperatorMetricGroup chainedOperatorMetricGroup = new InterceptingOperatorMetricGroup();
    InterceptingTaskMetricGroup taskMetricGroup = new InterceptingTaskMetricGroup() {

        @Override
        public InternalOperatorMetricGroup getOrAddOperator(OperatorID id, String name) {
            if (id.equals(mainOperatorId)) {
                return mainOperatorMetricGroup;
            } else if (id.equals(chainedOperatorId)) {
                return chainedOperatorMetricGroup;
            } else {
                return super.getOrAddOperator(id, name);
            }
        }
    };
    try (StreamTaskMailboxTestHarness<String> testHarness = new StreamTaskMailboxTestHarnessBuilder<>(MultipleInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO).modifyExecutionConfig(applyObjectReuse(objectReuse)).addInput(BasicTypeInfo.STRING_TYPE_INFO).addSourceInput(new SourceOperatorFactory<>(new MockSource(Boundedness.CONTINUOUS_UNBOUNDED, 2, true, false), WatermarkStrategy.forGenerator(ctx -> new RecordToWatermarkGenerator())), BasicTypeInfo.INT_TYPE_INFO).addInput(BasicTypeInfo.DOUBLE_TYPE_INFO).setupOperatorChain(mainOperatorId, new MapToStringMultipleInputOperatorFactory(3)).chain(chainedOperatorId, new WatermarkMetricOperator(), BasicTypeInfo.STRING_TYPE_INFO.createSerializer(new ExecutionConfig())).finish().setTaskMetricGroup(taskMetricGroup).build()) {
        Gauge<Long> taskInputWatermarkGauge = (Gauge<Long>) taskMetricGroup.get(MetricNames.IO_CURRENT_INPUT_WATERMARK);
        Gauge<Long> mainInput1WatermarkGauge = (Gauge<Long>) mainOperatorMetricGroup.get(MetricNames.currentInputWatermarkName(1));
        Gauge<Long> mainInput2WatermarkGauge = (Gauge<Long>) mainOperatorMetricGroup.get(MetricNames.currentInputWatermarkName(2));
        Gauge<Long> mainInput3WatermarkGauge = (Gauge<Long>) mainOperatorMetricGroup.get(MetricNames.currentInputWatermarkName(3));
        Gauge<Long> mainInputWatermarkGauge = (Gauge<Long>) mainOperatorMetricGroup.get(MetricNames.IO_CURRENT_INPUT_WATERMARK);
        Gauge<Long> mainOutputWatermarkGauge = (Gauge<Long>) mainOperatorMetricGroup.get(MetricNames.IO_CURRENT_OUTPUT_WATERMARK);
        Gauge<Long> chainedInputWatermarkGauge = (Gauge<Long>) chainedOperatorMetricGroup.get(MetricNames.IO_CURRENT_INPUT_WATERMARK);
        Gauge<Long> chainedOutputWatermarkGauge = (Gauge<Long>) chainedOperatorMetricGroup.get(MetricNames.IO_CURRENT_OUTPUT_WATERMARK);
        assertEquals(Long.MIN_VALUE, taskInputWatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainInputWatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainInput1WatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainInput2WatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainInput3WatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainOutputWatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, chainedInputWatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, chainedOutputWatermarkGauge.getValue().longValue());
        testHarness.processElement(new Watermark(1L), 0);
        assertEquals(Long.MIN_VALUE, taskInputWatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainInputWatermarkGauge.getValue().longValue());
        assertEquals(1L, mainInput1WatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainInput2WatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainInput3WatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainOutputWatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, chainedInputWatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, chainedOutputWatermarkGauge.getValue().longValue());
        addSourceRecords(testHarness, 1, 2);
        testHarness.processAll();
        assertEquals(Long.MIN_VALUE, taskInputWatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainInputWatermarkGauge.getValue().longValue());
        assertEquals(1L, mainInput1WatermarkGauge.getValue().longValue());
        assertEquals(2L, mainInput2WatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainInput3WatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, mainOutputWatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, chainedInputWatermarkGauge.getValue().longValue());
        assertEquals(Long.MIN_VALUE, chainedOutputWatermarkGauge.getValue().longValue());
        testHarness.processElement(new Watermark(2L), 1);
        assertEquals(1L, taskInputWatermarkGauge.getValue().longValue());
        assertEquals(1L, mainInputWatermarkGauge.getValue().longValue());
        assertEquals(1L, mainInput1WatermarkGauge.getValue().longValue());
        assertEquals(2L, mainInput2WatermarkGauge.getValue().longValue());
        assertEquals(2L, mainInput3WatermarkGauge.getValue().longValue());
        assertEquals(1L, mainOutputWatermarkGauge.getValue().longValue());
        assertEquals(1L, chainedInputWatermarkGauge.getValue().longValue());
        assertEquals(2L, chainedOutputWatermarkGauge.getValue().longValue());
        testHarness.processElement(new Watermark(4L), 0);
        addSourceRecords(testHarness, 1, 3);
        testHarness.processAll();
        assertEquals(2L, taskInputWatermarkGauge.getValue().longValue());
        assertEquals(2L, mainInputWatermarkGauge.getValue().longValue());
        assertEquals(4L, mainInput1WatermarkGauge.getValue().longValue());
        assertEquals(3L, mainInput2WatermarkGauge.getValue().longValue());
        assertEquals(2L, mainInput3WatermarkGauge.getValue().longValue());
        assertEquals(2L, mainOutputWatermarkGauge.getValue().longValue());
        assertEquals(2L, chainedInputWatermarkGauge.getValue().longValue());
        assertEquals(4L, chainedOutputWatermarkGauge.getValue().longValue());
        finishAddingRecords(testHarness, 1);
        testHarness.endInput();
        testHarness.waitForTaskCompletion();
        testHarness.finishProcessing();
    }
}

Example 3 with MockSource

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

the class MultipleInputStreamTaskTest method testTriggeringStopWithSavepointWithDrain.

@Test
public void testTriggeringStopWithSavepointWithDrain() throws Exception {
    SourceOperatorFactory<Integer> sourceOperatorFactory = new SourceOperatorFactory<>(new MockSource(Boundedness.CONTINUOUS_UNBOUNDED, 2), WatermarkStrategy.noWatermarks());
    CompletableFuture<Boolean> checkpointCompleted = new CompletableFuture<>();
    CheckpointResponder checkpointResponder = new TestCheckpointResponder() {

        @Override
        public void acknowledgeCheckpoint(JobID jobID, ExecutionAttemptID executionAttemptID, long checkpointId, CheckpointMetrics checkpointMetrics, TaskStateSnapshot subtaskState) {
            super.acknowledgeCheckpoint(jobID, executionAttemptID, checkpointId, checkpointMetrics, subtaskState);
            checkpointCompleted.complete(null);
        }
    };
    try (StreamTaskMailboxTestHarness<String> testHarness = new StreamTaskMailboxTestHarnessBuilder<>(MultipleInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO).setCollectNetworkEvents().modifyStreamConfig(config -> config.setCheckpointingEnabled(true)).modifyExecutionConfig(applyObjectReuse(objectReuse)).addInput(BasicTypeInfo.INT_TYPE_INFO).addInput(BasicTypeInfo.INT_TYPE_INFO).addInput(BasicTypeInfo.INT_TYPE_INFO).setTaskStateSnapshot(1, TaskStateSnapshot.FINISHED_ON_RESTORE).setupOperatorChain(new LifeCycleMonitorMultipleInputOperatorFactory()).finishForSingletonOperatorChain(StringSerializer.INSTANCE).setCheckpointResponder(checkpointResponder).build()) {
        CompletableFuture<Boolean> triggerResult = testHarness.streamTask.triggerCheckpointAsync(new CheckpointMetaData(2, 2), CheckpointOptions.alignedNoTimeout(SavepointType.terminate(SavepointFormatType.CANONICAL), CheckpointStorageLocationReference.getDefault()));
        checkpointCompleted.whenComplete((ignored, exception) -> testHarness.streamTask.notifyCheckpointCompleteAsync(2));
        testHarness.waitForTaskCompletion();
        testHarness.finishProcessing();
        assertTrue(triggerResult.isDone());
        assertTrue(triggerResult.get());
        assertTrue(checkpointCompleted.isDone());
    }
}

Example 4 with MockSource

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

the class MultipleInputStreamTaskTest method testCopyForObjectReuse.

@Test
public void testCopyForObjectReuse() throws Exception {
    SharedReference<List<Integer>> copiedElementsRef = sharedObjects.add(new ArrayList<>());
    CopyProxySerializer proxySerializer = new CopyProxySerializer(copiedElementsRef);
    try (StreamTaskMailboxTestHarness<String> testHarness = new StreamTaskMailboxTestHarnessBuilder<>(MultipleInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO).modifyExecutionConfig(applyObjectReuse(objectReuse)).addInput(BasicTypeInfo.STRING_TYPE_INFO).addSourceInput(new SourceOperatorFactory<>(new MockSource(Boundedness.BOUNDED, 1), WatermarkStrategy.noWatermarks()), proxySerializer).addInput(BasicTypeInfo.DOUBLE_TYPE_INFO).setupOutputForSingletonOperatorChain(new MapToStringMultipleInputOperatorFactory(3)).build()) {
        addSourceRecords(testHarness, 1, 42, 43);
        testHarness.endInput();
        testHarness.waitForTaskCompletion();
        if (objectReuse) {
            assertTrue(copiedElementsRef.get().isEmpty());
        } else {
            assertThat(copiedElementsRef.get(), containsInAnyOrder(42, 43));
        }
    }
}

Example 5 with MockSource

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

the class StreamGraphGeneratorExecutionModeDetectionTest method testExecutionModePropagationFromEnvWithDefaultAndUnboundedSource.

@Test
public void testExecutionModePropagationFromEnvWithDefaultAndUnboundedSource() {
    final StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
    environment.fromSource(new MockSource(Boundedness.CONTINUOUS_UNBOUNDED, 100), WatermarkStrategy.noWatermarks(), "unbounded-source").print();
    assertThat(environment.getStreamGraph(), hasProperties(GlobalStreamExchangeMode.ALL_EDGES_PIPELINED, JobType.STREAMING, false, true));
}