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

public AbstractStreamOperatorTestHarness(
		StreamOperator<OUT> operator,
		int maxParallelism,
		int parallelism,
		int subtaskIndex,
		OperatorID operatorID) throws Exception {
	this(
		operator,
		new MockEnvironmentBuilder()
			.setTaskName("MockTask")
			.setMemorySize(3 * 1024 * 1024)
			.setInputSplitProvider(new MockInputSplitProvider())
			.setBufferSize(1024)
			.setMaxParallelism(maxParallelism)
			.setParallelism(parallelism)
			.setSubtaskIndex(subtaskIndex)
			.build(),
		true,
		operatorID);
}
 

public AbstractStreamOperatorTestHarness(
		StreamOperatorFactory<OUT> factory,
		int maxParallelism,
		int parallelism,
		int subtaskIndex,
		OperatorID operatorID) throws Exception {
	this(
			null,
			factory,
			new MockEnvironmentBuilder()
					.setTaskName("MockTask")
					.setManagedMemorySize(3 * 1024 * 1024)
					.setInputSplitProvider(new MockInputSplitProvider())
					.setBufferSize(1024)
					.setMaxParallelism(maxParallelism)
					.setParallelism(parallelism)
					.setSubtaskIndex(subtaskIndex)
					.build(),
			true,
			operatorID);
}
 

public AbstractStreamOperatorTestHarness(
		StreamOperator<OUT> operator,
		int maxParallelism,
		int parallelism,
		int subtaskIndex,
		OperatorID operatorID) throws Exception {
	this(
			operator,
			SimpleOperatorFactory.of(operator),
			new MockEnvironmentBuilder()
					.setTaskName("MockTask")
					.setManagedMemorySize(3 * 1024 * 1024)
					.setInputSplitProvider(new MockInputSplitProvider())
					.setBufferSize(1024)
					.setMaxParallelism(maxParallelism)
					.setParallelism(parallelism)
					.setSubtaskIndex(subtaskIndex)
					.build(),
			true,
			operatorID);
}
 

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

/**
 * Tests that some StreamTask methods are called only in the main task's thread.
 * Currently, the main task's thread is the thread that creates the task.
 */
@Test
public void testThreadInvariants() throws Throwable {
	Configuration taskConfiguration = new Configuration();
	StreamConfig streamConfig = new StreamConfig(taskConfiguration);
	streamConfig.setStreamOperator(new StreamMap<>(value -> value));
	streamConfig.setOperatorID(new OperatorID());
	try (MockEnvironment mockEnvironment =
			new MockEnvironmentBuilder()
				.setTaskConfiguration(taskConfiguration)
				.build()) {

		ClassLoader taskClassLoader = new TestUserCodeClassLoader();

		RunningTask<ThreadInspectingTask> runningTask = runTask(() -> {
			Thread.currentThread().setContextClassLoader(taskClassLoader);
			return new ThreadInspectingTask(mockEnvironment);
		});
		runningTask.invocationFuture.get();

		assertThat(runningTask.streamTask.getTaskClassLoader(), is(sameInstance(taskClassLoader)));
	}
}
 

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

private static <T> SourceOperator<T, MockSourceSplit> createTestOperator(
		SourceReader<T, MockSourceSplit> reader,
		WatermarkStrategy<T> watermarkStrategy,
		ProcessingTimeService timeService) throws Exception {

	final OperatorStateStore operatorStateStore =
			new MemoryStateBackend().createOperatorStateBackend(
					new MockEnvironmentBuilder().build(),
					"test-operator",
					Collections.emptyList(),
					new CloseableRegistry());

	final StateInitializationContext stateContext = new StateInitializationContextImpl(
		false, operatorStateStore, null, null, null);

	final SourceOperator<T, MockSourceSplit> sourceOperator =
			new TestingSourceOperator<>(reader, watermarkStrategy, timeService);
	sourceOperator.initializeState(stateContext);
	sourceOperator.open();

	return sourceOperator;
}
 

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

public TestRuntimeContext(
	boolean isCheckpointingEnabled,
	int numParallelSubtasks,
	int subtaskIndex) {

	super(
		new TestStreamOperator(),
		new MockEnvironmentBuilder()
			.setTaskName("mockTask")
			.setManagedMemorySize(4 * MemoryManager.DEFAULT_PAGE_SIZE)
			.build(),
		Collections.emptyMap());

	this.isCheckpointingEnabled = isCheckpointingEnabled;
	this.numParallelSubtasks = numParallelSubtasks;
	this.subtaskIndex = subtaskIndex;
}
 

public TestRuntimeContext(
	boolean isCheckpointingEnabled,
	int numParallelSubtasks,
	int subtaskIndex) {

	super(
		new TestStreamOperator(),
		new MockEnvironmentBuilder()
			.setTaskName("mockTask")
			.setMemorySize(4 * MemoryManager.DEFAULT_PAGE_SIZE)
			.build(),
		Collections.emptyMap());

	this.isCheckpointingEnabled = isCheckpointingEnabled;
	this.numParallelSubtasks = numParallelSubtasks;
	this.subtaskIndex = subtaskIndex;
}
 

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

public MockStreamingRuntimeContext(
	boolean isCheckpointingEnabled,
	int numParallelSubtasks,
	int subtaskIndex) {

	super(
		new MockStreamOperator(),
		new MockEnvironmentBuilder()
			.setTaskName("mockTask")
			.setMemorySize(4 * MemoryManager.DEFAULT_PAGE_SIZE)
			.build(),
		Collections.emptyMap());

	this.isCheckpointingEnabled = isCheckpointingEnabled;
	this.numParallelSubtasks = numParallelSubtasks;
	this.subtaskIndex = subtaskIndex;
}
 

/**
 * Test set user code ClassLoader before calling ProcessingTimeCallback.
 */
@Test
public void testSetsUserCodeClassLoaderForTimerThreadFactory() throws Throwable {
	syncLatch = new OneShotLatch();

	try (MockEnvironment mockEnvironment =
		new MockEnvironmentBuilder()
			.setUserCodeClassLoader(new TestUserCodeClassLoader())
			.build()) {
		RunningTask<TimeServiceTask> task = runTask(() -> new TimeServiceTask(mockEnvironment));
		task.waitForTaskCompletion(false);

		assertThat(task.streamTask.getClassLoaders(), hasSize(greaterThanOrEqualTo(1)));
		assertThat(task.streamTask.getClassLoaders(), everyItem(instanceOf(TestUserCodeClassLoader.class)));
	}
}
 

public TestRuntimeContext(
	boolean isCheckpointingEnabled,
	int numParallelSubtasks,
	int subtaskIndex) {

	super(
		new TestStreamOperator(),
		new MockEnvironmentBuilder()
			.setTaskName("mockTask")
			.setMemorySize(4 * MemoryManager.DEFAULT_PAGE_SIZE)
			.build(),
		Collections.emptyMap());

	this.isCheckpointingEnabled = isCheckpointingEnabled;
	this.numParallelSubtasks = numParallelSubtasks;
	this.subtaskIndex = subtaskIndex;
}
 

public AbstractStreamOperatorTestHarness(
		StreamOperator<OUT> operator,
		int maxParallelism,
		int parallelism,
		int subtaskIndex,
		OperatorID operatorID) throws Exception {
	this(
		operator,
		new MockEnvironmentBuilder()
			.setTaskName("MockTask")
			.setMemorySize(3 * 1024 * 1024)
			.setInputSplitProvider(new MockInputSplitProvider())
			.setBufferSize(1024)
			.setMaxParallelism(maxParallelism)
			.setParallelism(parallelism)
			.setSubtaskIndex(subtaskIndex)
			.build(),
		true,
		operatorID);
}
 

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

public MockStreamingRuntimeContext(
	boolean isCheckpointingEnabled,
	int numParallelSubtasks,
	int subtaskIndex) {

	super(
		new MockStreamOperator(),
		new MockEnvironmentBuilder()
			.setTaskName("mockTask")
			.setMemorySize(4 * MemoryManager.DEFAULT_PAGE_SIZE)
			.build(),
		Collections.emptyMap());

	this.isCheckpointingEnabled = isCheckpointingEnabled;
	this.numParallelSubtasks = numParallelSubtasks;
	this.subtaskIndex = subtaskIndex;
}
 

/**
 * 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")
				.setManagedMemorySize(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 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);
		}
	}
 

@Test
public void testBeforeInvokeWithoutChannelStates() throws Exception {
	int numWriters = 2;
	int numGates = 2;
	RecoveryResultPartition[] partitions = new RecoveryResultPartition[numWriters];
	for (int i = 0; i < numWriters; i++) {
		partitions[i] = new RecoveryResultPartition();
	}
	RecoveryInputGate[] gates = new RecoveryInputGate[numGates];
	for (int i = 0; i < numGates; i++) {
		gates[i] = new RecoveryInputGate(partitions);
	}

	MockEnvironment mockEnvironment = new MockEnvironmentBuilder().build();
	mockEnvironment.addOutputs(asList(partitions));
	mockEnvironment.addInputs(asList(gates));
	StreamTask task = new MockStreamTaskBuilder(mockEnvironment).build();
	try {
		verifyResults(gates, partitions, false, false);

		task.beforeInvoke();

		verifyResults(gates, partitions, false, true);
	} finally {
		task.cleanUpInvoke();
	}
}
 

private OperatorStateStore createOperatorStateStore() throws Exception {
	MockEnvironment env = new MockEnvironmentBuilder().build();
	final AbstractStateBackend abstractStateBackend = new MemoryStateBackend();
	CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
	return abstractStateBackend.createOperatorStateBackend(
			env, "test-operator", Collections.emptyList(), cancelStreamRegistry);
}
 

/**
 * 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 RocksDBStateBackend backend = new RocksDBStateBackend(tempFolder.newFolder().toURI().toString());
	backend.setPriorityQueueStateType(RocksDBStateBackend.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 RocksDBStateBackend 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();
}
 

/**
 * Tests to cover case when user misuse optimizeForPointLookup with iterator interfaces on map state.
 *
 * <p>The option {@link ColumnFamilyOptions#optimizeForPointLookup(long)} would lead to iterator.seek with prefix bytes invalid.
 */
@Test
public void testMisuseOptimizePointLookupWithMapState() throws Exception {
	RocksDBStateBackend rocksDBStateBackend = createStateBackendWithOptimizePointLookup();
	RocksDBKeyedStateBackend<Integer> keyedStateBackend =
		createKeyedStateBackend(rocksDBStateBackend, new MockEnvironmentBuilder().build(), IntSerializer.INSTANCE);
	try {
		MapStateDescriptor<Integer, Long> stateDescriptor = new MapStateDescriptor<>("map", IntSerializer.INSTANCE, LongSerializer.INSTANCE);
		MapState<Integer, Long> mapState = keyedStateBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor);

		keyedStateBackend.setCurrentKey(1);
		Map<Integer, Long> expectedResult = new HashMap<>();
		for (int i = 0; i < 100; i++) {
			long uv = ThreadLocalRandom.current().nextLong();
			mapState.put(i, uv);
			expectedResult.put(i, uv);
		}

		Iterator<Map.Entry<Integer, Long>> iterator = mapState.entries().iterator();
		while (iterator.hasNext()) {
			Map.Entry<Integer, Long> entry = iterator.next();
			assertEquals(entry.getValue(), expectedResult.remove(entry.getKey()));
			iterator.remove();
		}
		assertTrue(expectedResult.isEmpty());
		assertTrue(mapState.isEmpty());
	} finally {
		keyedStateBackend.dispose();
	}
}
 

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

private MockEnvironment createMockEnvironment(String taskName) {
	return new MockEnvironmentBuilder()
		.setTaskName(taskName)
		.setManagedMemorySize(3 * 1024 * 1024)
		.setInputSplitProvider(new MockInputSplitProvider())
		.setBufferSize(1024)
		.build();
}
 

/**
 * Test set user code ClassLoader before calling ProcessingTimeCallback.
 */
@Test
public void testSetsUserCodeClassLoaderForTimerThreadFactory() throws Throwable {
	syncLatch = new OneShotLatch();

	try (MockEnvironment mockEnvironment =
		new MockEnvironmentBuilder()
			.setUserCodeClassLoader(new TestUserCodeClassLoader())
			.build()) {
		TimeServiceTask timerServiceTask = new TimeServiceTask(mockEnvironment);

		CompletableFuture<Void> invokeFuture = CompletableFuture.runAsync(
			() -> {
				try {
					timerServiceTask.invoke();
				} catch (Exception e) {
					throw new CompletionException(e);
				}
			},
			TestingUtils.defaultExecutor());

		invokeFuture.get();

		assertThat(timerServiceTask.getClassLoaders(), hasSize(greaterThanOrEqualTo(1)));
		assertThat(timerServiceTask.getClassLoaders(), everyItem(instanceOf(TestUserCodeClassLoader.class)));
	}
}
 

private MockEnvironment createMockEnvironment(String taskName) {
	return new MockEnvironmentBuilder()
		.setTaskName(taskName)
		.setMemorySize(3 * 1024 * 1024)
		.setInputSplitProvider(new MockInputSplitProvider())
		.setBufferSize(1024)
		.build();
}
 

private MockEnvironment createMockEnvironment(String taskName) {
	return new MockEnvironmentBuilder()
		.setTaskName(taskName)
		.setMemorySize(3 * 1024 * 1024)
		.setInputSplitProvider(new MockInputSplitProvider())
		.setBufferSize(1024)
		.build();
}
 

public MockStreamingRuntimeContext(
	boolean isCheckpointingEnabled,
	int numParallelSubtasks,
	int subtaskIndex) {

	this(isCheckpointingEnabled, numParallelSubtasks, subtaskIndex, new MockEnvironmentBuilder()
		.setTaskName("mockTask")
		.setManagedMemorySize(4 * MemoryManager.DEFAULT_PAGE_SIZE)
		.build());
}
 

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