Examples with IteratorWrappingTestSingleInputGate org.apache.flink.runtime.io.network.partition.consumer.IteratorWrappingTestSingleInputGate used on opensource projects

Example 1 with IteratorWrappingTestSingleInputGate

use of org.apache.flink.runtime.io.network.partition.consumer.IteratorWrappingTestSingleInputGate in project flink by apache.

the class DataSinkTaskTest method testUnionDataSinkTask.

@Test
public void testUnionDataSinkTask() {
    int keyCnt = 10;
    int valCnt = 20;
    super.initEnvironment(MEMORY_MANAGER_SIZE, NETWORK_BUFFER_SIZE);
    final IteratorWrappingTestSingleInputGate<?>[] readers = new IteratorWrappingTestSingleInputGate[4];
    readers[0] = super.addInput(new UniformRecordGenerator(keyCnt, valCnt, 0, 0, false), 0, false);
    readers[1] = super.addInput(new UniformRecordGenerator(keyCnt, valCnt, keyCnt, 0, false), 0, false);
    readers[2] = super.addInput(new UniformRecordGenerator(keyCnt, valCnt, keyCnt * 2, 0, false), 0, false);
    readers[3] = super.addInput(new UniformRecordGenerator(keyCnt, valCnt, keyCnt * 3, 0, false), 0, false);
    DataSinkTask<Record> testTask = new DataSinkTask<>(this.mockEnv);
    File tempTestFile = new File(tempFolder.getRoot(), UUID.randomUUID().toString());
    super.registerFileOutputTask(MockOutputFormat.class, tempTestFile.toURI().toString(), new Configuration());
    try {
        // which checks forwards existing notifications on registerListener calls.
        for (IteratorWrappingTestSingleInputGate<?> reader : readers) {
            reader.notifyNonEmpty();
        }
        testTask.invoke();
    } catch (Exception e) {
        LOG.debug("Exception while invoking the test task.", e);
        Assert.fail("Invoke method caused exception.");
    }
    Assert.assertTrue("Temp output file does not exist", tempTestFile.exists());
    FileReader fr = null;
    BufferedReader br = null;
    try {
        fr = new FileReader(tempTestFile);
        br = new BufferedReader(fr);
        HashMap<Integer, HashSet<Integer>> keyValueCountMap = new HashMap<>(keyCnt);
        while (br.ready()) {
            String line = br.readLine();
            Integer key = Integer.parseInt(line.substring(0, line.indexOf("_")));
            Integer val = Integer.parseInt(line.substring(line.indexOf("_") + 1, line.length()));
            if (!keyValueCountMap.containsKey(key)) {
                keyValueCountMap.put(key, new HashSet<Integer>());
            }
            keyValueCountMap.get(key).add(val);
        }
        Assert.assertTrue("Invalid key count in out file. Expected: " + keyCnt + " Actual: " + keyValueCountMap.keySet().size(), keyValueCountMap.keySet().size() == keyCnt * 4);
        for (Integer key : keyValueCountMap.keySet()) {
            Assert.assertTrue("Invalid value count for key: " + key + ". Expected: " + valCnt + " Actual: " + keyValueCountMap.get(key).size(), keyValueCountMap.get(key).size() == valCnt);
        }
    } catch (FileNotFoundException e) {
        Assert.fail("Out file got lost...");
    } catch (IOException ioe) {
        Assert.fail("Caught IOE while reading out file");
    } finally {
        if (br != null) {
            try {
                br.close();
            } catch (Throwable t) {
            }
        }
        if (fr != null) {
            try {
                fr.close();
            } catch (Throwable t) {
            }
        }
    }
}