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

Example 46 with MockEnvironment

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

the class SourceFunctionUtil method runRichSourceFunction.

private static <T extends Serializable> List<T> runRichSourceFunction(SourceFunction<T> sourceFunction) throws Exception {
    try (MockEnvironment environment = new MockEnvironmentBuilder().setTaskName("MockTask").setManagedMemorySize(3 * 1024 * 1024).setInputSplitProvider(new MockInputSplitProvider()).setBufferSize(1024).build()) {
        AbstractStreamOperator<?> operator = mock(AbstractStreamOperator.class);
        when(operator.getExecutionConfig()).thenReturn(new ExecutionConfig());
        RuntimeContext runtimeContext = new StreamingRuntimeContext(operator, environment, new HashMap<>());
        ((RichFunction) sourceFunction).setRuntimeContext(runtimeContext);
        ((RichFunction) sourceFunction).open(new Configuration());
        return runNonRichSourceFunction(sourceFunction);
    }
}

Example 47 with MockEnvironment

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

the class StateBackendTestBase method buildMockEnv.

private MockEnvironment buildMockEnv() throws Exception {
    MockEnvironment mockEnvironment = MockEnvironment.builder().setTaskStateManager(getTestTaskStateManager()).build();
    mockEnvironment.setCheckpointStorageAccess(getCheckpointStorageAccess());
    return mockEnvironment;
}

Example 48 with MockEnvironment

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

the class RocksDBAsyncSnapshotTest method testCleanupOfSnapshotsInFailureCase.

/**
 * Test that the snapshot files are cleaned up in case of a failure during the snapshot
 * procedure.
 */
@Test
public void testCleanupOfSnapshotsInFailureCase() throws Exception {
    long checkpointId = 1L;
    long timestamp = 42L;
    MockEnvironment env = MockEnvironment.builder().build();
    final IOException testException = new IOException("Test exception");
    CheckpointStateOutputStream outputStream = spy(new FailingStream(testException));
    RocksDBStateBackend backend = new RocksDBStateBackend((StateBackend) new MemoryStateBackend());
    backend.setDbStoragePath(temporaryFolder.newFolder().toURI().toString());
    AbstractKeyedStateBackend<Void> keyedStateBackend = backend.createKeyedStateBackend(env, new JobID(), "test operator", VoidSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), null, TtlTimeProvider.DEFAULT, new UnregisteredMetricsGroup(), Collections.emptyList(), new CloseableRegistry());
    try {
        // register a state so that the state backend has to checkpoint something
        keyedStateBackend.getPartitionedState("namespace", StringSerializer.INSTANCE, new ValueStateDescriptor<>("foobar", String.class));
        RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotFuture = keyedStateBackend.snapshot(checkpointId, timestamp, new TestCheckpointStreamFactory(() -> outputStream), CheckpointOptions.forCheckpointWithDefaultLocation());
        try {
            FutureUtils.runIfNotDoneAndGet(snapshotFuture);
            fail("Expected an exception to be thrown here.");
        } catch (ExecutionException e) {
            Assert.assertEquals(testException, e.getCause());
        }
        verify(outputStream).close();
    } finally {
        IOUtils.closeQuietly(keyedStateBackend);
        keyedStateBackend.dispose();
        IOUtils.closeQuietly(env);
    }
}

Example 49 with MockEnvironment

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

the class RocksDBStateBackendConfigTest method testConfigureTimerServiceLoadingFromApplication.

/**
 * Validates that user custom configuration from code should override the flink-conf.yaml.
 */
@Test
public void testConfigureTimerServiceLoadingFromApplication() throws Exception {
    final MockEnvironment env = new MockEnvironmentBuilder().build();
    // priorityQueueStateType of the job backend
    final EmbeddedRocksDBStateBackend backend = new EmbeddedRocksDBStateBackend();
    backend.setPriorityQueueStateType(EmbeddedRocksDBStateBackend.PriorityQueueStateType.HEAP);
    // priorityQueueStateType in the cluster config
    final Configuration configFromConfFile = new Configuration();
    configFromConfFile.setString(RocksDBOptions.TIMER_SERVICE_FACTORY.key(), RocksDBStateBackend.PriorityQueueStateType.ROCKSDB.toString());
    // configure final backend from job and cluster config
    final EmbeddedRocksDBStateBackend configuredRocksDBStateBackend = backend.configure(configFromConfFile, Thread.currentThread().getContextClassLoader());
    final RocksDBKeyedStateBackend<Integer> keyedBackend = createKeyedStateBackend(configuredRocksDBStateBackend, env, IntSerializer.INSTANCE);
    // priorityQueueStateType of the job backend should be preserved
    assertThat(keyedBackend.getPriorityQueueFactory(), instanceOf(HeapPriorityQueueSetFactory.class));
    keyedBackend.close();
    keyedBackend.dispose();
    env.close();
}

Example 50 with MockEnvironment

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

the class RocksDBStateBackendConfigTest method testConfigureTimerService.

@Test
public void testConfigureTimerService() throws Exception {
    final MockEnvironment env = getMockEnvironment(tempFolder.newFolder());
    // Fix the option key string
    Assert.assertEquals("state.backend.rocksdb.timer-service.factory", RocksDBOptions.TIMER_SERVICE_FACTORY.key());
    // Fix the option value string and ensure all are covered
    Assert.assertEquals(2, EmbeddedRocksDBStateBackend.PriorityQueueStateType.values().length);
    Assert.assertEquals("ROCKSDB", EmbeddedRocksDBStateBackend.PriorityQueueStateType.ROCKSDB.toString());
    Assert.assertEquals("HEAP", EmbeddedRocksDBStateBackend.PriorityQueueStateType.HEAP.toString());
    // Fix the default
    Assert.assertEquals(EmbeddedRocksDBStateBackend.PriorityQueueStateType.ROCKSDB, RocksDBOptions.TIMER_SERVICE_FACTORY.defaultValue());
    EmbeddedRocksDBStateBackend rocksDbBackend = new EmbeddedRocksDBStateBackend();
    RocksDBKeyedStateBackend<Integer> keyedBackend = createKeyedStateBackend(rocksDbBackend, env, IntSerializer.INSTANCE);
    Assert.assertEquals(RocksDBPriorityQueueSetFactory.class, keyedBackend.getPriorityQueueFactory().getClass());
    keyedBackend.dispose();
    Configuration conf = new Configuration();
    conf.set(RocksDBOptions.TIMER_SERVICE_FACTORY, EmbeddedRocksDBStateBackend.PriorityQueueStateType.HEAP);
    rocksDbBackend = rocksDbBackend.configure(conf, Thread.currentThread().getContextClassLoader());
    keyedBackend = createKeyedStateBackend(rocksDbBackend, env, IntSerializer.INSTANCE);
    Assert.assertEquals(HeapPriorityQueueSetFactory.class, keyedBackend.getPriorityQueueFactory().getClass());
    keyedBackend.dispose();
    env.close();
}