org.apache.flink.runtime.operators.testutils.MockEnvironment Java Examples

@Test
public void testNotifyCheckpointAbortedManyTimes() throws Exception {
	MockEnvironment mockEnvironment = MockEnvironment.builder().build();
	int maxRecordAbortedCheckpoints = 256;
	SubtaskCheckpointCoordinatorImpl subtaskCheckpointCoordinator = (SubtaskCheckpointCoordinatorImpl) new MockSubtaskCheckpointCoordinatorBuilder()
		.setEnvironment(mockEnvironment)
		.setMaxRecordAbortedCheckpoints(maxRecordAbortedCheckpoints)
		.build();

	final OperatorChain<?, ?> operatorChain = getOperatorChain(mockEnvironment);

	long notifyAbortedTimes = maxRecordAbortedCheckpoints + 42;
	for (int i = 1; i < notifyAbortedTimes; i++) {
		subtaskCheckpointCoordinator.notifyCheckpointAborted(i, operatorChain, () -> true);
		assertEquals(Math.min(maxRecordAbortedCheckpoints, i), subtaskCheckpointCoordinator.getAbortedCheckpointSize());
	}
}
 

@Test
public void testExecuteMailboxActionsAfterLeavingInputProcessorMailboxLoop() throws Exception {
	OneShotLatch latch = new OneShotLatch();
	try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder().build()) {
		RunningTask<StreamTask<?, ?>> task = runTask(() -> new StreamTask<Object, StreamOperator<Object>>(mockEnvironment) {
			@Override
			protected void init() throws Exception {
			}

			@Override
			protected void processInput(MailboxDefaultAction.Controller controller) throws Exception {
				mailboxProcessor.getMailboxExecutor(0).execute(latch::trigger, "trigger");
				controller.allActionsCompleted();
			}
		});
		latch.await();
		task.waitForTaskCompletion(false);
	}
}
 

/**
 * This test checks the async exceptions handling wraps the message and cause as an AsynchronousException
 * and propagates this to the environment.
 */
@Test
public void streamTaskAsyncExceptionHandler_handleException_forwardsMessageProperly() {
	MockEnvironment mockEnvironment = MockEnvironment.builder().build();
	RuntimeException expectedException = new RuntimeException("RUNTIME EXCEPTION");

	final StreamTask.StreamTaskAsyncExceptionHandler asyncExceptionHandler = new StreamTask.StreamTaskAsyncExceptionHandler(mockEnvironment);

	mockEnvironment.setExpectedExternalFailureCause(AsynchronousException.class);
	final String expectedErrorMessage = "EXPECTED_ERROR MESSAGE";

	asyncExceptionHandler.handleAsyncException(expectedErrorMessage, expectedException);

	// expect an AsynchronousException containing the supplied error details
	Optional<? extends Throwable> actualExternalFailureCause = mockEnvironment.getActualExternalFailureCause();
	final Throwable actualException = actualExternalFailureCause
		.orElseThrow(() -> new AssertionError("Expected exceptional completion"));

	assertThat(actualException, instanceOf(AsynchronousException.class));
	assertThat(actualException.getMessage(), is("EXPECTED_ERROR MESSAGE"));
	assertThat(actualException.getCause(), is(expectedException));
}
 

private static <T extends Serializable> List<T> runRichSourceFunction(SourceFunction<T> sourceFunction) throws Exception {
	try (MockEnvironment environment =
			new MockEnvironmentBuilder()
				.setTaskName("MockTask")
				.setMemorySize(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);
	}
}
 

/**
 * Verifies that latency metrics can be disabled via the {@link ExecutionConfig} even if they are enabled via
 * the configuration.
 */
@Test
public void testLatencyMarkEmissionDisabledOverrideViaExecutionConfig() throws Exception {
	testLatencyMarkEmission(0, (operator, timeProvider) -> {
		Configuration tmConfig = new Configuration();
		tmConfig.setLong(MetricOptions.LATENCY_INTERVAL, latencyMarkInterval);

		Environment env = MockEnvironment.builder()
			.setTaskManagerRuntimeInfo(new TestingTaskManagerRuntimeInfo(tmConfig))
			.build();

		ExecutionConfig executionConfig = new ExecutionConfig();
		executionConfig.setLatencyTrackingInterval(0);

		setupSourceOperator(operator, executionConfig, env, timeProvider);
	});
}
 

/**
 * Verifies that latency metrics can be enabled via the {@link ExecutionConfig} even if they are disabled via
 * the configuration.
 */
@Test
public void testLatencyMarkEmissionEnabledOverrideViaExecutionConfig() throws Exception {
	testLatencyMarkEmission((int) (maxProcessingTime / latencyMarkInterval) + 1, (operator, timeProvider) -> {
		ExecutionConfig executionConfig = new ExecutionConfig();
		executionConfig.setLatencyTrackingInterval(latencyMarkInterval);

		Configuration tmConfig = new Configuration();
		tmConfig.setLong(MetricOptions.LATENCY_INTERVAL, 0L);

		Environment env = MockEnvironment.builder()
			.setTaskManagerRuntimeInfo(new TestingTaskManagerRuntimeInfo(tmConfig))
			.build();

		setupSourceOperator(operator, executionConfig, env, timeProvider);
	});
}
 

SubtaskCheckpointCoordinator build() throws IOException {
	if (environment == null) {
		this.environment = MockEnvironment.builder().build();
	}
	if (checkpointStorage == null) {
		this.checkpointStorage = new MemoryBackendCheckpointStorage(environment.getJobID(), null, null, 1024);
	}
	if (asyncExceptionHandler == null) {
		this.asyncExceptionHandler = new NonHandleAsyncException();
	}

	return new SubtaskCheckpointCoordinatorImpl(
		checkpointStorage,
		taskName,
		actionExecutor,
		closeableRegistry,
		executorService,
		environment,
		asyncExceptionHandler,
		unalignedCheckpointEnabled,
		prepareInputSnapshot,
		maxRecordAbortedCheckpoints);
}
 

/**
 * Verifies that latency metrics can be disabled via the {@link ExecutionConfig} even if they are enabled via
 * the configuration.
 */
@Test
public void testLatencyMarkEmissionDisabledOverrideViaExecutionConfig() throws Exception {
	testLatencyMarkEmission(0, (operator, timeProvider) -> {
		Configuration tmConfig = new Configuration();
		tmConfig.setLong(MetricOptions.LATENCY_INTERVAL, latencyMarkInterval);

		Environment env = MockEnvironment.builder()
			.setTaskManagerRuntimeInfo(new TestingTaskManagerRuntimeInfo(tmConfig))
			.build();

		ExecutionConfig executionConfig = new ExecutionConfig();
		executionConfig.setLatencyTrackingInterval(0);

		setupSourceOperator(operator, executionConfig, env, timeProvider);
	});
}
 

@Test
public void testProcessWithUnAvailableOutput() throws Exception {
	try (final MockEnvironment environment = setupEnvironment(new boolean[] {true, false})) {
		final int numberOfProcessCalls = 10;
		final AvailabilityTestInputProcessor inputProcessor = new AvailabilityTestInputProcessor(numberOfProcessCalls);
		final StreamTask task = new MockStreamTaskBuilder(environment)
			.setStreamInputProcessor(inputProcessor)
			.build();
		final MailboxExecutor executor = task.mailboxProcessor.getMainMailboxExecutor();

		final RunnableWithException completeFutureTask = () -> {
			assertEquals(1, inputProcessor.currentNumProcessCalls);
			assertTrue(task.mailboxProcessor.isDefaultActionUnavailable());
			environment.getWriter(1).getAvailableFuture().complete(null);
		};

		executor.submit(() -> {
			executor.submit(completeFutureTask, "This task will complete the future to resume process input action."); },
			"This task will submit another task to execute after processing input once.");

		task.invoke();
		assertEquals(numberOfProcessCalls, inputProcessor.currentNumProcessCalls);
	}
}
 

/**
 * Verifies that latency metrics can be enabled via the {@link ExecutionConfig} even if they are disabled via
 * the configuration.
 */
@Test
public void testLatencyMarkEmissionEnabledOverrideViaExecutionConfig() throws Exception {
	testLatencyMarkEmission((int) (maxProcessingTime / latencyMarkInterval) + 1, (operator, timeProvider) -> {
		ExecutionConfig executionConfig = new ExecutionConfig();
		executionConfig.setLatencyTrackingInterval(latencyMarkInterval);

		Configuration tmConfig = new Configuration();
		tmConfig.setLong(MetricOptions.LATENCY_INTERVAL, 0L);

		Environment env = MockEnvironment.builder()
			.setTaskManagerRuntimeInfo(new TestingTaskManagerRuntimeInfo(tmConfig))
			.build();

		setupSourceOperator(operator, executionConfig, env, timeProvider);
	});
}
 

@Before
public void before() throws Exception {
	ioManager = new IOManagerAsync();
	MockEnvironment environment = new MockEnvironmentBuilder()
		.setTaskName("mockTask")
		.setManagedMemorySize(4 * MemoryManager.DEFAULT_PAGE_SIZE)
		.setIOManager(ioManager)
		.build();
	runtimeContext = new MockStreamingRuntimeContext(false, 1, 0, environment);

	gateway = new MockOperatorEventGateway();
	coordinator = new CollectSinkOperatorCoordinator(SOCKET_TIMEOUT_MILLIS);
	coordinator.start();

	// only used in checkpointed tests
	functionInitializationContext = new MockFunctionInitializationContext();

	jobFinished = false;
}
 

/**
 * Verifies that latency metrics can be disabled via the {@link ExecutionConfig} even if they are enabled via
 * the configuration.
 */
@Test
public void testLatencyMarkEmissionDisabledOverrideViaExecutionConfig() throws Exception {
	testLatencyMarkEmission(0, (operator, timeProvider) -> {
		Configuration tmConfig = new Configuration();
		tmConfig.setLong(MetricOptions.LATENCY_INTERVAL, latencyMarkInterval);

		Environment env = MockEnvironment.builder()
			.setTaskManagerRuntimeInfo(new TestingTaskManagerRuntimeInfo(tmConfig))
			.build();

		ExecutionConfig executionConfig = new ExecutionConfig();
		executionConfig.setLatencyTrackingInterval(0);

		setupSourceOperator(operator, executionConfig, env, timeProvider);
	});
}
 

@Test
public void testNotifyCheckpointComplete() throws Exception {
	TestTaskStateManager stateManager = new TestTaskStateManager();
	MockEnvironment mockEnvironment = MockEnvironment.builder().setTaskStateManager(stateManager).build();
	SubtaskCheckpointCoordinator subtaskCheckpointCoordinator = new MockSubtaskCheckpointCoordinatorBuilder()
		.setEnvironment(mockEnvironment)
		.build();

	final OperatorChain<?, ?> operatorChain = getOperatorChain(mockEnvironment);

	long checkpointId = 42L;
	{
		subtaskCheckpointCoordinator.notifyCheckpointComplete(checkpointId, operatorChain, () -> true);
		assertEquals(checkpointId, stateManager.getNotifiedCompletedCheckpointId());
	}

	long newCheckpointId = checkpointId + 1;
	{
		subtaskCheckpointCoordinator.notifyCheckpointComplete(newCheckpointId, operatorChain, () -> false);
		// even task is not running, state manager could still receive the notification.
		assertEquals(newCheckpointId, stateManager.getNotifiedCompletedCheckpointId());
	}
}
 

@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);
	}
}
 

private StreamingRuntimeContext createRuntimeContext(
		AtomicReference<Object> descriptorCapture,
		ExecutionConfig config) throws Exception {
	return createDescriptorCapturingMockOp(
		descriptorCapture,
		config,
		MockEnvironment.builder()
			.setExecutionConfig(config)
			.build()).getRuntimeContext();
}
 

private StreamingRuntimeContext createRuntimeContext(AbstractStreamOperator<?> operator) {
	return new StreamingRuntimeContext(
		MockEnvironment.builder()
			.build(),
		Collections.emptyMap(),
		operator.getMetricGroup(),
		operator.getOperatorID(),
		operator.getProcessingTimeService(),
		operator.getKeyedStateStore(),
		ExternalResourceInfoProvider.NO_EXTERNAL_RESOURCES);
}
 

public OneInputStreamOperatorTestHarness(
	OneInputStreamOperator<IN, OUT> operator,
	TypeSerializer<IN> typeSerializerIn,
	MockEnvironment environment) throws Exception {
	this(operator, environment);

	config.setTypeSerializersIn(Preconditions.checkNotNull(typeSerializerIn));
}
 

/**
 * Verifies that latency metrics can be enabled via the {@link ExecutionConfig}.
 */
@Test
public void testLatencyMarkEmissionEnabledViaExecutionConfig() throws Exception {
	testLatencyMarkEmission((int) (maxProcessingTime / latencyMarkInterval) + 1, (operator, timeProvider) -> {
		ExecutionConfig executionConfig = new ExecutionConfig();
		executionConfig.setLatencyTrackingInterval(latencyMarkInterval);

		setupSourceOperator(operator, executionConfig, MockEnvironment.builder().build(), timeProvider);
	});
}
 

@Test
public void testSavepointNotResultingInPriorityEvents() throws Exception {
	MockEnvironment mockEnvironment = MockEnvironment.builder().build();

	SubtaskCheckpointCoordinator coordinator = new MockSubtaskCheckpointCoordinatorBuilder()
			.setUnalignedCheckpointEnabled(true)
			.setEnvironment(mockEnvironment)
			.build();

	AtomicReference<Boolean> broadcastedPriorityEvent = new AtomicReference<>(null);
	final OperatorChain<?, ?> operatorChain = new OperatorChain(
			new MockStreamTaskBuilder(mockEnvironment).build(),
			new NonRecordWriter<>()) {
		@Override
		public void broadcastEvent(AbstractEvent event, boolean isPriorityEvent) throws IOException {
			super.broadcastEvent(event, isPriorityEvent);
			broadcastedPriorityEvent.set(isPriorityEvent);
		}
	};

	coordinator.checkpointState(
			new CheckpointMetaData(0, 0),
			new CheckpointOptions(SAVEPOINT, CheckpointStorageLocationReference.getDefault()),
			new CheckpointMetrics(),
			operatorChain,
			() -> false);

	assertEquals(false, broadcastedPriorityEvent.get());
}
 

private StreamingRuntimeContext createRuntimeContext() throws Exception {
	return new StreamingRuntimeContext(
		createListPlainMockOp(),
		MockEnvironment.builder()
			.build(),
		Collections.emptyMap());
}
 

/**
 * This test checks the behavior for basic setting of local DB directories.
 */
@Test
public void testSetDbPath() throws Exception {
	final RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI().toString());

	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();
	}
}
 

/**
 * Tests that the StreamTask first closes all of its operators before setting its
 * state to not running (isRunning == false)
 *
 * <p>See FLINK-7430.
 */
@Test
public void testOperatorClosingBeforeStopRunning() throws Throwable {
	BlockingCloseStreamOperator.resetLatches();
	Configuration taskConfiguration = new Configuration();
	StreamConfig streamConfig = new StreamConfig(taskConfiguration);
	streamConfig.setStreamOperator(new BlockingCloseStreamOperator());
	streamConfig.setOperatorID(new OperatorID());

	try (MockEnvironment mockEnvironment =
			new MockEnvironmentBuilder()
				.setTaskName("Test Task")
				.setMemorySize(32L * 1024L)
				.setInputSplitProvider(new MockInputSplitProvider())
				.setBufferSize(1)
				.setTaskConfiguration(taskConfiguration)
				.build()) {

		RunningTask<StreamTask<Void, BlockingCloseStreamOperator>> task = runTask(() -> new NoOpStreamTask<>(mockEnvironment));

		BlockingCloseStreamOperator.inClose.await();

		// check that the StreamTask is not yet in isRunning == false
		assertTrue(task.streamTask.isRunning());

		// let the operator finish its close operation
		BlockingCloseStreamOperator.finishClose.trigger();

		task.waitForTaskCompletion(false);

		// now the StreamTask should no longer be running
		assertFalse(task.streamTask.isRunning());
	}
}
 

private MockEnvironment setupEnvironment(boolean[] outputAvailabilities) {
	final Configuration configuration = new Configuration();
	new MockStreamConfig(configuration, outputAvailabilities.length);

	final List<ResultPartitionWriter> writers = new ArrayList<>(outputAvailabilities.length);
	for (int i = 0; i < outputAvailabilities.length; i++) {
		writers.add(new AvailabilityTestResultPartitionWriter(outputAvailabilities[i]));
	}

	final MockEnvironment environment = new MockEnvironmentBuilder()
		.setTaskConfiguration(configuration)
		.build();
	environment.addOutputs(writers);
	return environment;
}
 

public OneInputStreamOperatorTestHarness(
	OneInputStreamOperator<IN, OUT> operator,
	MockEnvironment environment) throws Exception {
	super(operator, environment);

	this.oneInputOperator = operator;
}
 

public OneInputStreamOperatorTestHarness(
	OneInputStreamOperator<IN, OUT> operator,
	TypeSerializer<IN> typeSerializerIn,
	MockEnvironment environment) throws Exception {
	this(operator, environment);

	config.setTypeSerializerIn1(Preconditions.checkNotNull(typeSerializerIn));
}
 

/**
 * Verifies that by default no latency metrics are emitted.
 */
@Test
public void testLatencyMarkEmissionDisabled() throws Exception {
	testLatencyMarkEmission(0, (operator, timeProvider) -> {
		setupSourceOperator(operator, new ExecutionConfig(), MockEnvironment.builder().build(), timeProvider);
	});
}
 

/**
 * Verifies that latency metrics can be enabled via the {@link ExecutionConfig}.
 */
@Test
public void testLatencyMarkEmissionEnabledViaExecutionConfig() throws Exception {
	testLatencyMarkEmission((int) (maxProcessingTime / latencyMarkInterval) + 1, (operator, timeProvider) -> {
		ExecutionConfig executionConfig = new ExecutionConfig();
		executionConfig.setLatencyTrackingInterval(latencyMarkInterval);

		setupSourceOperator(operator, executionConfig, MockEnvironment.builder().build(), timeProvider);
	});
}
 

public KeyedOneInputStreamOperatorTestHarness(
		final OneInputStreamOperator<IN, OUT> operator,
		final  KeySelector<IN, K> keySelector,
		final TypeInformation<K> keyType,
		final MockEnvironment environment) throws Exception {

	super(operator, environment);

	ClosureCleaner.clean(keySelector, ExecutionConfig.ClosureCleanerLevel.RECURSIVE, false);
	config.setStatePartitioner(0, keySelector);
	config.setStateKeySerializer(keyType.createSerializer(executionConfig));
}
 

@Test
public void testProcessWithAvailableOutput() throws Exception {
	try (final MockEnvironment environment = setupEnvironment(new boolean[] {true, true})) {
		final int numberOfProcessCalls = 10;
		final AvailabilityTestInputProcessor inputProcessor = new AvailabilityTestInputProcessor(numberOfProcessCalls);
		final StreamTask task = new MockStreamTaskBuilder(environment)
			.setStreamInputProcessor(inputProcessor)
			.build();

		task.invoke();
		assertEquals(numberOfProcessCalls, inputProcessor.currentNumProcessCalls);
	}
}
 

public KeyedOneInputStreamOperatorTestHarness(
		final OneInputStreamOperator<IN, OUT> operator,
		final KeySelector<IN, K> keySelector,
		final TypeInformation<K> keyType,
		final MockEnvironment environment) throws Exception {
	super(operator, environment);

	ClosureCleaner.clean(keySelector, ExecutionConfig.ClosureCleanerLevel.RECURSIVE, false);
	config.setStatePartitioner(0, keySelector);
	config.setStateKeySerializer(keyType.createSerializer(executionConfig));
}