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

Example 11 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment in project flink by apache.

the class InputFormatSourceFunctionTest method testFormatLifecycle.

private void testFormatLifecycle(final boolean midCancel) throws Exception {
    final int noOfSplits = 5;
    final int cancelAt = 2;
    final LifeCycleTestInputFormat format = new LifeCycleTestInputFormat();
    final InputFormatSourceFunction<Integer> reader = new InputFormatSourceFunction<>(format, TypeInformation.of(Integer.class));
    try (MockEnvironment environment = new MockEnvironmentBuilder().setTaskName("no").setManagedMemorySize(4 * MemoryManager.DEFAULT_PAGE_SIZE).build()) {
        reader.setRuntimeContext(new MockRuntimeContext(format, noOfSplits, environment));
        Assert.assertTrue(!format.isConfigured);
        Assert.assertTrue(!format.isInputFormatOpen);
        Assert.assertTrue(!format.isSplitOpen);
        reader.open(new Configuration());
        Assert.assertTrue(format.isConfigured);
        TestSourceContext ctx = new TestSourceContext(reader, format, midCancel, cancelAt);
        reader.run(ctx);
        int splitsSeen = ctx.getSplitsSeen();
        Assert.assertTrue(midCancel ? splitsSeen == cancelAt : splitsSeen == noOfSplits);
        // we have exhausted the splits so the
        // format and splits should be closed by now
        Assert.assertTrue(!format.isSplitOpen);
        Assert.assertTrue(!format.isInputFormatOpen);
    }
}

Example 12 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment in project flink by apache.

the class RocksDBStateBackendConfigTest method testFailWhenNoLocalStorageDir.

// ------------------------------------------------------------------------
// RocksDB local file directory initialization
// ------------------------------------------------------------------------
@Test
public void testFailWhenNoLocalStorageDir() throws Exception {
    final File targetDir = tempFolder.newFolder();
    Assume.assumeTrue("Cannot mark directory non-writable", targetDir.setWritable(false, false));
    String checkpointPath = tempFolder.newFolder().toURI().toString();
    RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
    try (MockEnvironment env = getMockEnvironment(tempFolder.newFolder())) {
        rocksDbBackend.setDbStoragePath(targetDir.getAbsolutePath());
        boolean hasFailure = false;
        try {
            rocksDbBackend.createKeyedStateBackend(env, env.getJobID(), "foobar", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()), TtlTimeProvider.DEFAULT, new UnregisteredMetricsGroup(), Collections.emptyList(), new CloseableRegistry());
        } catch (Exception e) {
            assertTrue(e.getMessage().contains("No local storage directories available"));
            assertTrue(e.getMessage().contains(targetDir.getAbsolutePath()));
            hasFailure = true;
        }
        assertTrue("We must see a failure because no storaged directory is feasible.", hasFailure);
    } finally {
        // noinspection ResultOfMethodCallIgnored
        targetDir.setWritable(true, false);
    }
}

Example 13 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment in project flink by apache.

the class RocksDBStateBackendConfigTest method testContinueOnSomeDbDirectoriesMissing.

@Test
public void testContinueOnSomeDbDirectoriesMissing() throws Exception {
    final File targetDir1 = tempFolder.newFolder();
    final File targetDir2 = tempFolder.newFolder();
    Assume.assumeTrue("Cannot mark directory non-writable", targetDir1.setWritable(false, false));
    String checkpointPath = tempFolder.newFolder().toURI().toString();
    RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
    try (MockEnvironment env = getMockEnvironment(tempFolder.newFolder())) {
        rocksDbBackend.setDbStoragePaths(targetDir1.getAbsolutePath(), targetDir2.getAbsolutePath());
        try {
            AbstractKeyedStateBackend<Integer> keyedStateBackend = rocksDbBackend.createKeyedStateBackend(env, env.getJobID(), "foobar", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()), TtlTimeProvider.DEFAULT, new UnregisteredMetricsGroup(), Collections.emptyList(), new CloseableRegistry());
            IOUtils.closeQuietly(keyedStateBackend);
            keyedStateBackend.dispose();
        } catch (Exception e) {
            e.printStackTrace();
            fail("Backend initialization failed even though some paths were available");
        }
    } finally {
        // noinspection ResultOfMethodCallIgnored
        targetDir1.setWritable(true, false);
    }
}

Example 14 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment in project flink by apache.

the class RocksDBStateBackendConfigTest method testSetDbPath.

// ------------------------------------------------------------------------
// RocksDB local file directory
// ------------------------------------------------------------------------
/**
 * This test checks the behavior for basic setting of local DB directories.
 */
@Test
public void testSetDbPath() throws Exception {
    final EmbeddedRocksDBStateBackend rocksDbBackend = new EmbeddedRocksDBStateBackend();
    final String testDir1 = tempFolder.newFolder().getAbsolutePath();
    final String testDir2 = tempFolder.newFolder().getAbsolutePath();
    assertNull(rocksDbBackend.getDbStoragePaths());
    rocksDbBackend.setDbStoragePath(testDir1);
    assertArrayEquals(new String[] { testDir1 }, rocksDbBackend.getDbStoragePaths());
    rocksDbBackend.setDbStoragePath(null);
    assertNull(rocksDbBackend.getDbStoragePaths());
    rocksDbBackend.setDbStoragePaths(testDir1, testDir2);
    assertArrayEquals(new String[] { testDir1, testDir2 }, rocksDbBackend.getDbStoragePaths());
    final MockEnvironment env = getMockEnvironment(tempFolder.newFolder());
    final RocksDBKeyedStateBackend<Integer> keyedBackend = createKeyedStateBackend(rocksDbBackend, env, IntSerializer.INSTANCE);
    try {
        File instanceBasePath = keyedBackend.getInstanceBasePath();
        assertThat(instanceBasePath.getAbsolutePath(), anyOf(startsWith(testDir1), startsWith(testDir2)));
        // noinspection NullArgumentToVariableArgMethod
        rocksDbBackend.setDbStoragePaths(null);
        assertNull(rocksDbBackend.getDbStoragePaths());
    } finally {
        IOUtils.closeQuietly(keyedBackend);
        keyedBackend.dispose();
        env.close();
    }
}

Example 15 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment in project flink by apache.

the class StreamOperatorChainingTest method testMultiChainingWithSplit.

/**
 * Verify that multi-chaining works with object reuse enabled.
 */
private void testMultiChainingWithSplit(StreamExecutionEnvironment env) throws Exception {
    // set parallelism to 2 to avoid chaining with source in case when available processors is
    // 1.
    env.setParallelism(2);
    // the actual elements will not be used
    DataStream<Integer> input = env.fromElements(1, 2, 3);
    sink1Results = new ArrayList<>();
    sink2Results = new ArrayList<>();
    sink3Results = new ArrayList<>();
    input = input.map(value -> value);
    OutputTag<Integer> oneOutput = new OutputTag<Integer>("one") {
    };
    OutputTag<Integer> otherOutput = new OutputTag<Integer>("other") {
    };
    SingleOutputStreamOperator<Object> split = input.process(new ProcessFunction<Integer, Object>() {

        private static final long serialVersionUID = 1L;

        @Override
        public void processElement(Integer value, Context ctx, Collector<Object> out) throws Exception {
            if (value.equals(1)) {
                ctx.output(oneOutput, value);
            } else {
                ctx.output(otherOutput, value);
            }
        }
    });
    split.getSideOutput(oneOutput).map(value -> "First 1: " + value).addSink(new SinkFunction<String>() {

        @Override
        public void invoke(String value, Context ctx) throws Exception {
            sink1Results.add(value);
        }
    });
    split.getSideOutput(oneOutput).map(value -> "First 2: " + value).addSink(new SinkFunction<String>() {

        @Override
        public void invoke(String value, Context ctx) throws Exception {
            sink2Results.add(value);
        }
    });
    split.getSideOutput(otherOutput).map(value -> "Second: " + value).addSink(new SinkFunction<String>() {

        @Override
        public void invoke(String value, Context ctx) throws Exception {
            sink3Results.add(value);
        }
    });
    // be build our own StreamTask and OperatorChain
    JobGraph jobGraph = env.getStreamGraph().getJobGraph();
    Assert.assertTrue(jobGraph.getVerticesSortedTopologicallyFromSources().size() == 2);
    JobVertex chainedVertex = jobGraph.getVerticesSortedTopologicallyFromSources().get(1);
    Configuration configuration = chainedVertex.getConfiguration();
    StreamConfig streamConfig = new StreamConfig(configuration);
    StreamMap<Integer, Integer> headOperator = streamConfig.getStreamOperator(Thread.currentThread().getContextClassLoader());
    try (MockEnvironment environment = createMockEnvironment(chainedVertex.getName())) {
        StreamTask<Integer, StreamMap<Integer, Integer>> mockTask = createMockTask(streamConfig, environment);
        OperatorChain<Integer, StreamMap<Integer, Integer>> operatorChain = createOperatorChain(streamConfig, environment, mockTask);
        headOperator.setup(mockTask, streamConfig, operatorChain.getMainOperatorOutput());
        operatorChain.initializeStateAndOpenOperators(null);
        headOperator.processElement(new StreamRecord<>(1));
        headOperator.processElement(new StreamRecord<>(2));
        headOperator.processElement(new StreamRecord<>(3));
        assertThat(sink1Results, contains("First 1: 1"));
        assertThat(sink2Results, contains("First 2: 1"));
        assertThat(sink3Results, contains("Second: 2", "Second: 3"));
    }
}