org.apache.flink.runtime.testutils.CommonTestUtils Java Examples

public static void verifyStateRestore(
		JobVertexID jobVertexID, ExecutionJobVertex executionJobVertex,
		List<KeyGroupRange> keyGroupPartitions) throws Exception {

	for (int i = 0; i < executionJobVertex.getParallelism(); i++) {

		JobManagerTaskRestore taskRestore = executionJobVertex.getTaskVertices()[i].getCurrentExecutionAttempt().getTaskRestore();
		Assert.assertEquals(1L, taskRestore.getRestoreCheckpointId());
		TaskStateSnapshot stateSnapshot = taskRestore.getTaskStateSnapshot();

		OperatorSubtaskState operatorState = stateSnapshot.getSubtaskStateByOperatorID(OperatorID.fromJobVertexID(jobVertexID));

		ChainedStateHandle<OperatorStateHandle> expectedOpStateBackend =
				generateChainedPartitionableStateHandle(jobVertexID, i, 2, 8, false);

		assertTrue(CommonTestUtils.isSteamContentEqual(
				expectedOpStateBackend.get(0).openInputStream(),
				operatorState.getManagedOperatorState().iterator().next().openInputStream()));

		KeyGroupsStateHandle expectPartitionedKeyGroupState = generateKeyGroupState(
				jobVertexID, keyGroupPartitions.get(i), false);
		compareKeyedState(Collections.singletonList(expectPartitionedKeyGroupState), operatorState.getManagedKeyedState());
	}
}
 

public static void verifyStateRestore(
		JobVertexID jobVertexID, ExecutionJobVertex executionJobVertex,
		List<KeyGroupRange> keyGroupPartitions) throws Exception {

	for (int i = 0; i < executionJobVertex.getParallelism(); i++) {

		JobManagerTaskRestore taskRestore = executionJobVertex.getTaskVertices()[i].getCurrentExecutionAttempt().getTaskRestore();
		Assert.assertEquals(1L, taskRestore.getRestoreCheckpointId());
		TaskStateSnapshot stateSnapshot = taskRestore.getTaskStateSnapshot();

		OperatorSubtaskState operatorState = stateSnapshot.getSubtaskStateByOperatorID(OperatorID.fromJobVertexID(jobVertexID));

		ChainedStateHandle<OperatorStateHandle> expectedOpStateBackend =
				generateChainedPartitionableStateHandle(jobVertexID, i, 2, 8, false);

		assertTrue(CommonTestUtils.isStreamContentEqual(
				expectedOpStateBackend.get(0).openInputStream(),
				operatorState.getManagedOperatorState().iterator().next().openInputStream()));

		KeyGroupsStateHandle expectPartitionedKeyGroupState = generateKeyGroupState(
				jobVertexID, keyGroupPartitions.get(i), false);
		compareKeyedState(Collections.singletonList(expectPartitionedKeyGroupState), operatorState.getManagedKeyedState());
	}
}
 

public static void verifyStateRestore(
	JobVertexID jobVertexID, ExecutionJobVertex executionJobVertex,
	List<KeyGroupRange> keyGroupPartitions) throws Exception {

	for (int i = 0; i < executionJobVertex.getParallelism(); i++) {

		JobManagerTaskRestore taskRestore = executionJobVertex.getTaskVertices()[i].getCurrentExecutionAttempt().getTaskRestore();
		Assert.assertEquals(1L, taskRestore.getRestoreCheckpointId());
		TaskStateSnapshot stateSnapshot = taskRestore.getTaskStateSnapshot();

		OperatorSubtaskState operatorState = stateSnapshot.getSubtaskStateByOperatorID(OperatorID.fromJobVertexID(jobVertexID));

		ChainedStateHandle<OperatorStateHandle> expectedOpStateBackend =
			generateChainedPartitionableStateHandle(jobVertexID, i, 2, 8, false);

		assertTrue(CommonTestUtils.isStreamContentEqual(
			expectedOpStateBackend.get(0).openInputStream(),
			operatorState.getManagedOperatorState().iterator().next().openInputStream()));

		KeyGroupsStateHandle expectPartitionedKeyGroupState = generateKeyGroupState(
			jobVertexID, keyGroupPartitions.get(i), false);
		compareKeyedState(Collections.singletonList(expectPartitionedKeyGroupState), operatorState.getManagedKeyedState());
	}
}
 

private static void waitUntilJobIsRestarted(
	final RestClusterClient<ApplicationId> restClusterClient,
	final JobID jobId,
	final int expectedFullRestarts) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> getJobFullRestarts(restClusterClient, jobId) >= expectedFullRestarts,
		Deadline.fromNow(TIMEOUT));
}
 

private static void waitUntilJobIsRunning(RestClusterClient<ApplicationId> restClusterClient, JobID jobId) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> {
			final JobDetailsInfo jobDetails = restClusterClient.getJobDetails(jobId).get();
			return jobDetails.getJobVertexInfos()
				.stream()
				.map(toExecutionState())
				.allMatch(isRunning());
		},
		Deadline.fromNow(TIMEOUT));
}
 

private void killApplicationAndWait(final ApplicationId id) throws Exception {
	final YarnClient yarnClient = getYarnClient();
	checkState(yarnClient != null, "yarnClient must be initialized");

	yarnClient.killApplication(id);

	CommonTestUtils.waitUntilCondition(() -> !yarnClient.getApplications(EnumSet.of(YarnApplicationState.KILLED, YarnApplicationState.FINISHED)).isEmpty(),
		Deadline.fromNow(TIMEOUT));
}
 

private void waitForApplicationAttempt(final ApplicationId applicationId, final int attemptId) throws Exception {
	final YarnClient yarnClient = getYarnClient();
	checkState(yarnClient != null, "yarnClient must be initialized");

	CommonTestUtils.waitUntilCondition(() -> {
		final ApplicationReport applicationReport = yarnClient.getApplicationReport(applicationId);
		return applicationReport.getCurrentApplicationAttemptId().getAttemptId() >= attemptId;
	}, Deadline.fromNow(TIMEOUT));
}
 

private CompletableFuture<JobResult> submitJobAndWaitUntilRunning(JobGraph jobGraph) throws Exception {
	miniCluster.submitJob(jobGraph).get();

	final CompletableFuture<JobResult> jobResultFuture = miniCluster.requestJobResult(jobGraph.getJobID());

	assertThat(jobResultFuture.isDone(), is(false));

	CommonTestUtils.waitUntilCondition(
		jobIsRunning(() -> miniCluster.getExecutionGraph(jobGraph.getJobID())),
		Deadline.fromNow(TESTING_TIMEOUT),
		50L);

	return jobResultFuture;
}
 

private void waitUntilAllExecutionsAreScheduled(final JobMasterGateway jobMasterGateway) throws Exception {
	final Duration duration = Duration.ofMillis(testingTimeout.toMilliseconds());
	final Deadline deadline = Deadline.fromNow(duration);

	CommonTestUtils.waitUntilCondition(
		() -> getExecutions(jobMasterGateway).stream().allMatch(execution -> execution.getState() == ExecutionState.SCHEDULED),
		deadline);
}
 

public TestProcessBuilder(String mainClass) throws IOException {
	File tempLogFile = File.createTempFile(getClass().getSimpleName() + "-", "-log4j.properties");
	tempLogFile.deleteOnExit();
	CommonTestUtils.printLog4jDebugConfig(tempLogFile);

	jvmArgs.add("-Dlog.level=DEBUG");
	jvmArgs.add("-Dlog4j.configuration=file:" + tempLogFile.getAbsolutePath());
	jvmArgs.add("-classpath");
	jvmArgs.add(getCurrentClasspath());

	this.mainClass = mainClass;
}
 

@Test
public void testLoadingDynamicPropertiesInBash() throws Exception {
	final Configuration clientConfiguration = new Configuration();
	final File root = folder.getRoot();
	final File homeDir = new File(root, "home");
	assertTrue(homeDir.mkdir());
	BootstrapTools.writeConfiguration(clientConfiguration, new File(homeDir, FLINK_CONF_FILENAME));

	final Configuration jmUpdatedConfiguration = getJobManagerUpdatedConfiguration();

	final File shellScriptFile = generateLaunchContainerScript(
		homeDir,
		BootstrapTools.getDynamicPropertiesAsString(clientConfiguration, jmUpdatedConfiguration));

	Process process = new ProcessBuilder(shellScriptFile.getAbsolutePath()).start();
	try {
		final StringWriter processOutput = new StringWriter();
		new CommonTestUtils.PipeForwarder(process.getErrorStream(), processOutput);

		if (!process.waitFor(10, TimeUnit.SECONDS)) {
			throw new Exception("TestingTaskManagerRunner did not shutdown in time.");
		}
		assertEquals(processOutput.toString(), 0, process.exitValue());
	} finally {
		process.destroy();
	}
}
 

public TestProcessBuilder(String mainClass) throws IOException {
	File tempLogFile = File.createTempFile(getClass().getSimpleName() + "-", "-log4j.properties");
	tempLogFile.deleteOnExit();
	CommonTestUtils.printLog4jDebugConfig(tempLogFile);

	jvmArgs.add("-Dlog.level=DEBUG");
	jvmArgs.add("-Dlog4j.configurationFile=file:" + tempLogFile.getAbsolutePath());
	jvmArgs.add("-classpath");
	jvmArgs.add(getCurrentClasspath());

	this.mainClass = mainClass;
}
 

@Test
public void testCancelBatchResult() throws Exception {
	ResultDescriptor resultDescriptor = new ResultDescriptor("", schema, true, true);

	TestingExecutor mockExecutor = new TestingExecutorBuilder()
		.setSnapshotResultSupplier(TypedResult::empty)
		.build();

	CliTableauResultView view = new CliTableauResultView(
			terminal, mockExecutor, "session", resultDescriptor);

	// submit result display in another thread
	ExecutorService executorService = Executors.newSingleThreadExecutor();
	Future<?> furture = executorService.submit(view::displayBatchResults);

	// wait until we trying to get batch result
	CommonTestUtils.waitUntilCondition(
		() -> mockExecutor.getNumSnapshotResultCalls() > 0,
		Deadline.now().plus(Duration.ofSeconds(5)),
		50L);

	// send signal to cancel
	terminal.raise(Terminal.Signal.INT);
	furture.get(5, TimeUnit.SECONDS);

	Assert.assertEquals("Query terminated" + System.lineSeparator(), terminalOutput.toString());
	// didn't have a chance to read page
	assertThat(mockExecutor.getNumRetrieveResultPageCalls(), is(0));
	// tried to cancel query
	assertThat(mockExecutor.getNumCancelCalls(), is(1));

	view.close();
}
 

private static void assertNumberOfSlotsPerTask(
		final String host,
		final int port,
		final int slotsNumber) throws Exception {
	try {
		CommonTestUtils.waitUntilCondition(() -> getNumberOfSlotsPerTaskManager(host, port) == slotsNumber, Deadline.fromNow(Duration.ofSeconds(30)));
	} catch (final TimeoutException e) {
		final int currentNumberOfSlots = getNumberOfSlotsPerTaskManager(host, port);
		fail(String.format("Expected slots per TM to be %d, was: %d", slotsNumber, currentNumberOfSlots));
	}
}
 

private static void assertNumberOfSlotsPerTask(
		final String host,
		final int port,
		final int slotsNumber) throws Exception {
	try {
		CommonTestUtils.waitUntilCondition(() -> getNumberOfSlotsPerTaskManager(host, port) == slotsNumber, Deadline.fromNow(Duration.ofSeconds(30)));
	} catch (final TimeoutException e) {
		final int currentNumberOfSlots = getNumberOfSlotsPerTaskManager(host, port);
		fail(String.format("Expected slots per TM to be %d, was: %d", slotsNumber, currentNumberOfSlots));
	}
}
 

private static void waitUntilJobIsRestarted(
	final RestClusterClient<ApplicationId> restClusterClient,
	final JobID jobId,
	final int expectedFullRestarts) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> getJobFullRestarts(restClusterClient, jobId) >= expectedFullRestarts,
		Deadline.fromNow(TIMEOUT));
}
 

private static void waitUntilJobIsRunning(RestClusterClient<ApplicationId> restClusterClient, JobID jobId) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> {
			final JobDetailsInfo jobDetails = restClusterClient.getJobDetails(jobId).get();
			return jobDetails.getJobVertexInfos()
				.stream()
				.map(toExecutionState())
				.allMatch(isRunning());
		},
		Deadline.fromNow(TIMEOUT));
}
 

private void killApplicationAndWait(final ApplicationId id) throws Exception {
	final YarnClient yarnClient = getYarnClient();
	checkState(yarnClient != null, "yarnClient must be initialized");

	yarnClient.killApplication(id);

	CommonTestUtils.waitUntilCondition(() -> !yarnClient.getApplications(EnumSet.of(YarnApplicationState.KILLED, YarnApplicationState.FINISHED)).isEmpty(),
		Deadline.fromNow(TIMEOUT));
}
 

private void waitForApplicationAttempt(final ApplicationId applicationId, final int attemptId) throws Exception {
	final YarnClient yarnClient = getYarnClient();
	checkState(yarnClient != null, "yarnClient must be initialized");

	CommonTestUtils.waitUntilCondition(() -> {
		final ApplicationReport applicationReport = yarnClient.getApplicationReport(applicationId);
		return applicationReport.getCurrentApplicationAttemptId().getAttemptId() >= attemptId;
	}, Deadline.fromNow(TIMEOUT));
}
 

private CompletableFuture<JobResult> submitJobAndWaitUntilRunning(JobGraph jobGraph) throws Exception {
	miniCluster.submitJob(jobGraph).get();

	final CompletableFuture<JobResult> jobResultFuture = miniCluster.requestJobResult(jobGraph.getJobID());

	assertThat(jobResultFuture.isDone(), is(false));

	CommonTestUtils.waitUntilCondition(
		jobIsRunning(() -> miniCluster.getExecutionGraph(jobGraph.getJobID())),
		Deadline.fromNow(TESTING_TIMEOUT),
		50L);

	return jobResultFuture;
}
 

private void waitUntilAllExecutionsAreScheduled(final JobMasterGateway jobMasterGateway) throws Exception {
	final Duration duration = Duration.ofMillis(testingTimeout.toMilliseconds());
	final Deadline deadline = Deadline.fromNow(duration);

	CommonTestUtils.waitUntilCondition(
		() -> getExecutions(jobMasterGateway).stream().allMatch(execution -> execution.getState() == ExecutionState.SCHEDULED),
		deadline);
}
 

public TestProcessBuilder(String mainClass) throws IOException {
	File tempLogFile = File.createTempFile(getClass().getSimpleName() + "-", "-log4j.properties");
	tempLogFile.deleteOnExit();
	CommonTestUtils.printLog4jDebugConfig(tempLogFile);

	jvmArgs.add("-Dlog.level=DEBUG");
	jvmArgs.add("-Dlog4j.configuration=file:" + tempLogFile.getAbsolutePath());
	jvmArgs.add("-classpath");
	jvmArgs.add(getCurrentClasspath());

	this.mainClass = mainClass;
}
 

private static void assertNumberOfSlotsPerTask(
		final String host,
		final int port,
		final int slotsNumber) throws Exception {
	try {
		CommonTestUtils.waitUntilCondition(() -> getNumberOfSlotsPerTaskManager(host, port) == slotsNumber, Deadline.fromNow(Duration.ofSeconds(30)));
	} catch (final TimeoutException e) {
		final int currentNumberOfSlots = getNumberOfSlotsPerTaskManager(host, port);
		fail(String.format("Expected slots per TM to be %d, was: %d", slotsNumber, currentNumberOfSlots));
	}
}
 

private CompletableFuture<JobResult> submitJobAndWaitUntilRunning(JobGraph jobGraph) throws Exception {
	miniCluster.submitJob(jobGraph).get();

	final CompletableFuture<JobResult> jobResultFuture = miniCluster.requestJobResult(jobGraph.getJobID());

	assertThat(jobResultFuture.isDone(), is(false));

	CommonTestUtils.waitUntilCondition(
		jobIsRunning(() -> miniCluster.getExecutionGraph(jobGraph.getJobID())),
		Deadline.fromNow(TESTING_TIMEOUT),
		50L);

	return jobResultFuture;
}
 

private static void waitUntilJobIsRestarted(
	final RestClusterClient<ApplicationId> restClusterClient,
	final JobID jobId,
	final int expectedFullRestarts) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> getJobFullRestarts(restClusterClient, jobId) >= expectedFullRestarts,
		Deadline.fromNow(TIMEOUT));
}
 

private static void waitUntilJobIsRunning(RestClusterClient<ApplicationId> restClusterClient, JobID jobId) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> {
			final JobDetailsInfo jobDetails = restClusterClient.getJobDetails(jobId).get();
			return jobDetails.getJobVertexInfos()
				.stream()
				.map(toExecutionState())
				.allMatch(isRunning());
		},
		Deadline.fromNow(TIMEOUT));
}
 

private void waitUntilBackPressureStatsAvailable() {
	try {
		CommonTestUtils.waitUntilCondition(
			() -> {
				final Optional<OperatorBackPressureStats> stats = getBackPressureStats();
				return stats.isPresent();
				},
			Deadline.fromNow(TIMEOUT));
	} catch (Exception e) {
		throw new RuntimeException(e);
	}
}
 

private void killApplicationAndWait(final ApplicationId id) throws Exception {
	final YarnClient yarnClient = getYarnClient();
	checkState(yarnClient != null, "yarnClient must be initialized");

	yarnClient.killApplication(id);

	CommonTestUtils.waitUntilCondition(() -> !yarnClient.getApplications(EnumSet.of(YarnApplicationState.KILLED, YarnApplicationState.FINISHED)).isEmpty(),
		Deadline.fromNow(TIMEOUT));
}
 

private void waitForApplicationAttempt(final ApplicationId applicationId, final int attemptId) throws Exception {
	final YarnClient yarnClient = getYarnClient();
	checkState(yarnClient != null, "yarnClient must be initialized");

	CommonTestUtils.waitUntilCondition(() -> {
		final ApplicationReport applicationReport = yarnClient.getApplicationReport(applicationId);
		return applicationReport.getCurrentApplicationAttemptId().getAttemptId() >= attemptId;
	}, Deadline.fromNow(TIMEOUT));
}
 

private void waitUntilTaskExecutorsHaveConnected(int numTaskExecutors, Deadline deadline) throws Exception {
	CommonTestUtils.waitUntilCondition(
		() -> miniCluster.requestClusterOverview().get().getNumTaskManagersConnected() == numTaskExecutors,
		deadline,
		10L);
}