Examples with TestWithCountDownSink edu.snu.mist.core.task.utils.TestWithCountDownSink used on opensource projects

Example 1 with TestWithCountDownSink

use of edu.snu.mist.core.task.utils.TestWithCountDownSink in project mist by snuspl.

the class QueryManagerTest method testSubmitComplexQueryHelper.

@SuppressWarnings("unchecked")
private void testSubmitComplexQueryHelper(final Configuration conf) throws Exception {
    final String queryId = "testQuery";
    final List<String> inputs = Arrays.asList("mist is a cloud of tiny water droplets suspended in the atmosphere", "a mist rose out of the river", "the peaks were shrouded in mist");
    // Expected results
    final List<Map<String, Integer>> intermediateResult = getIntermediateResult(inputs);
    final List<String> expectedSink1Output = intermediateResult.stream().map(input -> input.toString()).collect(Collectors.toList());
    final List<Integer> expectedSink2Output = intermediateResult.stream().map(totalCountMapFunc).collect(Collectors.toList());
    // Number of expected outputs
    final CountDownLatch countDownAllOutputs = new CountDownLatch(intermediateResult.size() * 2);
    // Create the execution DAG of the query
    final ExecutionDag executionDag = new ExecutionDag(new AdjacentListDAG<>());
    // Create source
    final TestDataGenerator dataGenerator = new TestDataGenerator(inputs);
    final EventGenerator eventGenerator = new PunctuatedEventGenerator(null, input -> false, null, 0, null, null);
    final PhysicalSource src = new PhysicalSourceImpl("testSource", new HashMap<>(), dataGenerator, eventGenerator);
    final Injector injector = Tang.Factory.getTang().newInjector(conf);
    // Create sinks
    final List<String> sink1Result = new LinkedList<>();
    final List<Integer> sink2Result = new LinkedList<>();
    final PhysicalSink sink1 = new PhysicalSinkImpl<>("sink1", null, new TestWithCountDownSink<String>(sink1Result, countDownAllOutputs));
    final PhysicalSink sink2 = new PhysicalSinkImpl<>("sink2", null, new TestWithCountDownSink<Integer>(sink2Result, countDownAllOutputs));
    // Fake operator chain dag of QueryManager
    final AvroDag fakeAvroDag = new AvroDag();
    fakeAvroDag.setQueryId(queryId);
    // fakeAvroDag.setSuperGroupId("testGroup");
    // Construct execution dag
    constructExecutionDag(executionDag, src, sink1, sink2);
    // Create mock DagGenerator. It returns the above  execution dag
    final ConfigDagGenerator configDagGenerator = mock(ConfigDagGenerator.class);
    final DAG<ConfigVertex, MISTEdge> configDag = mock(DAG.class);
    when(configDagGenerator.generate(fakeAvroDag)).thenReturn(configDag);
    final DagGenerator dagGenerator = mock(DagGenerator.class);
    // when(dagGenerator.generate(configDag, tuple.getValue().getJarFilePaths())).thenReturn(executionDag);
    // Build QueryManager
    final QueryManager queryManager = queryManagerBuild(fakeAvroDag, configDagGenerator, dagGenerator, injector);
    queryManager.create(fakeAvroDag);
    // Wait until all of the outputs are generated
    countDownAllOutputs.await();
    // Check the outputs
    Assert.assertEquals(expectedSink1Output, sink1Result);
    Assert.assertEquals(expectedSink2Output, sink2Result);
    src.close();
    queryManager.close();
    // Delete plan directory and plans
    deletePlans(injector);
}