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

Example 11 with MockSourceSplit

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

the class SourceCoordinatorTestBase method addTestingSplitSet.

protected void addTestingSplitSet(int num) {
    final List<MockSourceSplit> splits = new ArrayList<>();
    for (int i = 0; i < num; i++) {
        splits.add(new MockSourceSplit(i));
    }
    getEnumerator().addNewSplits(splits);
}

Example 12 with MockSourceSplit

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

the class SourceOperatorStreamTaskTest method executeAndWaitForCheckpoint.

private TaskStateSnapshot executeAndWaitForCheckpoint(long checkpointId, TaskStateSnapshot initialSnapshot, IntStream expectedRecords) throws Exception {
    try (StreamTaskMailboxTestHarness<Integer> testHarness = createTestHarness(checkpointId, initialSnapshot)) {
        // Add records to the splits.
        MockSourceSplit split = getAndMaybeAssignSplit(testHarness);
        // Add records to the split and update expected output.
        addRecords(split, NUM_RECORDS);
        // Process all the records.
        testHarness.processAll();
        CheckpointOptions checkpointOptions = CheckpointOptions.forCheckpointWithDefaultLocation();
        triggerCheckpointWaitForFinish(testHarness, checkpointId, checkpointOptions);
        // Build expected output to verify the results
        Queue<Object> expectedOutput = new LinkedList<>();
        expectedRecords.forEach(r -> expectedOutput.offer(new StreamRecord<>(r, TimestampAssigner.NO_TIMESTAMP)));
        // Add barrier to the expected output.
        expectedOutput.add(new CheckpointBarrier(checkpointId, checkpointId, checkpointOptions));
        assertEquals(checkpointId, testHarness.taskStateManager.getReportedCheckpointId());
        assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
        return testHarness.taskStateManager.getLastJobManagerTaskStateSnapshot();
    }
}

Example 13 with MockSourceSplit

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

the class SourceOperatorEventTimeTest method testSequenceOfEvents.

@SuppressWarnings("FinalPrivateMethod")
@SafeVarargs
private final List<Object> testSequenceOfEvents(final boolean emitProgressiveWatermarks, final WatermarkStrategy<Integer> watermarkStrategy, final Consumer<ReaderOutput<Integer>>... actions) throws Exception {
    final CollectingDataOutput<Integer> out = new CollectingDataOutput<>();
    final TestProcessingTimeService timeService = new TestProcessingTimeService();
    // start somewhere that is not zero
    timeService.setCurrentTime(Integer.MAX_VALUE);
    final SourceReader<Integer, MockSourceSplit> reader = new InterpretingSourceReader(actions);
    final SourceOperator<Integer, MockSourceSplit> sourceOperator = createTestOperator(reader, watermarkStrategy, timeService, emitProgressiveWatermarks);
    while (sourceOperator.emitNext(out) != DataInputStatus.END_OF_INPUT) {
        timeService.setCurrentTime(timeService.getCurrentProcessingTime() + 100);
    }
    return out.events;
}

Example 14 with MockSourceSplit

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

the class SourceOperatorStreamTaskTest method getAndMaybeAssignSplit.

private MockSourceSplit getAndMaybeAssignSplit(StreamTaskMailboxTestHarness<Integer> testHarness) throws Exception {
    List<MockSourceSplit> assignedSplits = getSourceReaderFromTask(testHarness).getAssignedSplits();
    if (assignedSplits.isEmpty()) {
        // Prepare the source split and assign it to the source reader.
        MockSourceSplit split = new MockSourceSplit(0, 0);
        // Assign the split to the source reader.
        AddSplitEvent<MockSourceSplit> addSplitEvent = new AddSplitEvent<>(Collections.singletonList(split), new MockSourceSplitSerializer());
        testHarness.getStreamTask().dispatchOperatorEvent(OPERATOR_ID, new SerializedValue<>(addSplitEvent));
        // Run the task until the split assignment is done.
        while (assignedSplits.isEmpty()) {
            testHarness.getStreamTask().runMailboxStep();
        }
        // Need to mark the source reader as available for further processing.
        getSourceReaderFromTask(testHarness).markAvailable();
    }
    // The source reader already has an assigned split, just return it
    return assignedSplits.get(0);
}

Example 15 with MockSourceSplit

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

the class SourceCoordinatorContextTest method testAssignSplits.

@SuppressWarnings("unchecked")
private void testAssignSplits(boolean fromCoordinatorExecutor) throws Exception {
    sourceReady();
    registerReaders();
    // Assign splits to the readers.
    SplitsAssignment<MockSourceSplit> splitsAssignment = getSplitsAssignment(2, 0);
    if (fromCoordinatorExecutor) {
        coordinatorExecutor.submit(() -> context.assignSplits(splitsAssignment)).get();
    } else {
        context.assignSplits(splitsAssignment);
    }
    // The tracker should have recorded the assignments.
    verifyAssignment(Collections.singletonList("0"), splitSplitAssignmentTracker.uncheckpointedAssignments().get(0));
    verifyAssignment(Arrays.asList("1", "2"), splitSplitAssignmentTracker.uncheckpointedAssignments().get(1));
    // The OperatorCoordinatorContext should have received the event sending call.
    assertEquals("There should be two events sent to the subtasks.", 2, receivingTasks.getNumberOfSentEvents());
    // Assert the events to subtask0.
    List<OperatorEvent> eventsToSubtask0 = receivingTasks.getSentEventsForSubtask(0);
    assertEquals(1, eventsToSubtask0.size());
    OperatorEvent event = eventsToSubtask0.get(0);
    assertTrue(event instanceof AddSplitEvent);
    verifyAssignment(Collections.singletonList("0"), ((AddSplitEvent<MockSourceSplit>) event).splits(new MockSourceSplitSerializer()));
}