Examples with DataSourceTask org.apache.flink.runtime.operators.DataSourceTask used on opensource projects

Example 1 with DataSourceTask

use of org.apache.flink.runtime.operators.DataSourceTask in project flink by apache.

the class ChainTaskTest method testDataSourceTaskOutputInCloseMethod.

@Test
public void testDataSourceTaskOutputInCloseMethod() throws IOException {
    final int numChainedTasks = 10;
    final int keyCnt = 100;
    final int valCnt = 10;
    final File tempTestFile = new File(tempFolder.getRoot(), UUID.randomUUID().toString());
    DataSourceTaskTest.InputFilePreparator.prepareInputFile(new UniformRecordGenerator(keyCnt, valCnt, false), tempTestFile, true);
    initEnvironment(MEMORY_MANAGER_SIZE, NETWORK_BUFFER_SIZE);
    addOutput(outList);
    final DataSourceTask<Record> testTask = new DataSourceTask<>(mockEnv);
    registerFileInputTask(testTask, DataSourceTaskTest.MockInputFormat.class, tempTestFile.toURI().toString(), "\n");
    for (int i = 0; i < numChainedTasks; i++) {
        final TaskConfig taskConfig = new TaskConfig(new Configuration());
        taskConfig.addOutputShipStrategy(ShipStrategyType.FORWARD);
        taskConfig.setOutputSerializer(serFact);
        taskConfig.setStubWrapper(new UserCodeClassWrapper<>(ChainTaskTest.MockDuplicateLastValueMapFunction.class));
        getTaskConfig().addChainedTask(ChainedFlatMapDriver.class, taskConfig, "chained-" + i);
    }
    try {
        testTask.invoke();
        Assert.assertEquals(keyCnt * valCnt + numChainedTasks, outList.size());
    } catch (Exception e) {
        e.printStackTrace();
        Assert.fail("Invoke method caused exception.");
    }
}