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

Example 1 with MockSourceSplitSerializer

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

the class SourceOperatorTest method testHandleAddSplitsEvent.

@Test
public void testHandleAddSplitsEvent() throws Exception {
    operator.initializeState(context.createStateContext());
    operator.open();
    MockSourceSplit newSplit = new MockSourceSplit((2));
    operator.handleOperatorEvent(new AddSplitEvent<>(Collections.singletonList(newSplit), new MockSourceSplitSerializer()));
    // The source reader should have been assigned two splits.
    assertEquals(Arrays.asList(SourceOperatorTestContext.MOCK_SPLIT, newSplit), mockSourceReader.getAssignedSplits());
}

Example 2 with MockSourceSplitSerializer

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

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

Example 4 with MockSourceSplitSerializer

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

the class SourceCoordinatorContextTest method testCallableInterruptedDuringShutdownDoNotFailJob.

@Test
public void testCallableInterruptedDuringShutdownDoNotFailJob() throws InterruptedException {
    AtomicReference<Throwable> expectedError = new AtomicReference<>(null);
    ManuallyTriggeredScheduledExecutorService manualWorkerExecutor = new ManuallyTriggeredScheduledExecutorService();
    ManuallyTriggeredScheduledExecutorService manualCoordinatorExecutor = new ManuallyTriggeredScheduledExecutorService();
    SourceCoordinatorContext<MockSourceSplit> testingContext = new SourceCoordinatorContext<>(manualCoordinatorExecutor, manualWorkerExecutor, new SourceCoordinatorProvider.CoordinatorExecutorThreadFactory(TEST_OPERATOR_ID.toHexString(), getClass().getClassLoader()), operatorCoordinatorContext, new MockSourceSplitSerializer(), splitSplitAssignmentTracker);
    testingContext.callAsync(() -> {
        throw new InterruptedException();
    }, (ignored, e) -> {
        if (e != null) {
            expectedError.set(e);
            throw new RuntimeException(e);
        }
    });
    manualWorkerExecutor.triggerAll();
    testingContext.close();
    manualCoordinatorExecutor.triggerAll();
    assertTrue(expectedError.get() instanceof InterruptedException);
    assertFalse(operatorCoordinatorContext.isJobFailed());
}

Example 5 with MockSourceSplitSerializer

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

the class MultipleInputStreamTaskTest method addSourceRecords.

static void addSourceRecords(StreamTaskMailboxTestHarness<String> testHarness, int sourceId, Boundedness boundedness, int... records) throws Exception {
    OperatorID sourceOperatorID = getSourceOperatorID(testHarness, sourceId);
    // Prepare the source split and assign it to the source reader.
    MockSourceSplit split = new MockSourceSplit(0, 0, boundedness == Boundedness.BOUNDED ? records.length : Integer.MAX_VALUE);
    for (int record : records) {
        split.addRecord(record);
    }
    // Assign the split to the source reader.
    AddSplitEvent<MockSourceSplit> addSplitEvent = new AddSplitEvent<>(Collections.singletonList(split), new MockSourceSplitSerializer());
    testHarness.getStreamTask().dispatchOperatorEvent(sourceOperatorID, new SerializedValue<>(addSplitEvent));
}