Java Code Examples for org.apache.flink.runtime.webmonitor.TestingDispatcherGateway#Builder

@Test
public void testSuccessfulJobSubmission() throws Exception {
	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();
	try (ObjectOutputStream objectOut = new ObjectOutputStream(Files.newOutputStream(jobGraphFile))) {
		objectOut.writeObject(new JobGraph("testjob"));
	}

	TestingDispatcherGateway.Builder builder = new TestingDispatcherGateway.Builder();
	builder
		.setBlobServerPort(blobServer.getPort())
		.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
		.setHostname("localhost");
	DispatcherGateway mockGateway = builder.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.getFileName().toString(), Collections.emptyList(), Collections.emptyList());

	handler.handleRequest(new HandlerRequest<>(request, EmptyMessageParameters.getInstance(), Collections.emptyMap(), Collections.emptyMap(), Collections.singleton(jobGraphFile.toFile())), mockGateway)
		.get();
}
 

@Test
public void testSuccessfulJobSubmission() throws Exception {
	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();
	try (ObjectOutputStream objectOut = new ObjectOutputStream(Files.newOutputStream(jobGraphFile))) {
		objectOut.writeObject(new JobGraph("testjob"));
	}

	TestingDispatcherGateway.Builder builder = new TestingDispatcherGateway.Builder();
	builder
		.setBlobServerPort(blobServer.getPort())
		.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
		.setHostname("localhost");
	DispatcherGateway mockGateway = builder.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.getFileName().toString(), Collections.emptyList(), Collections.emptyList());

	handler.handleRequest(new HandlerRequest<>(request, EmptyMessageParameters.getInstance(), Collections.emptyMap(), Collections.emptyMap(), Collections.singleton(jobGraphFile.toFile())), mockGateway)
		.get();
}
 

@Test
public void testClusterShutdownWhenApplicationFails() throws Exception {
	// we're "listening" on this to be completed to verify that the cluster
	// is being shut down from the ApplicationDispatcherBootstrap
	final CompletableFuture<ApplicationStatus> externalShutdownFuture = new CompletableFuture<>();

	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
			.setRequestJobStatusFunction(jobId -> CompletableFuture.completedFuture(JobStatus.FAILED))
			.setRequestJobResultFunction(jobId -> CompletableFuture.completedFuture(createFailedJobResult(jobId)))
			.setClusterShutdownFunction((status) -> {
				externalShutdownFuture.complete(status);
				return CompletableFuture.completedFuture(Acknowledge.get());
			});

	ApplicationDispatcherBootstrap bootstrap = createApplicationDispatcherBootstrap(3);

	final CompletableFuture<Acknowledge> shutdownFuture =
			bootstrap.runApplicationAndShutdownClusterAsync(dispatcherBuilder.build(), scheduledExecutor);

	// wait until the bootstrap "thinks" it's done
	shutdownFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);

	// verify that the dispatcher is actually being shut down
	assertThat(externalShutdownFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS), is(ApplicationStatus.FAILED));
}
 

@Test
public void testSuccessfulJobSubmission() throws Exception {
	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();
	try (ObjectOutputStream objectOut = new ObjectOutputStream(Files.newOutputStream(jobGraphFile))) {
		objectOut.writeObject(new JobGraph("testjob"));
	}

	TestingDispatcherGateway.Builder builder = new TestingDispatcherGateway.Builder();
	builder
		.setBlobServerPort(blobServer.getPort())
		.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
		.setHostname("localhost");
	DispatcherGateway mockGateway = builder.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.getFileName().toString(), Collections.emptyList(), Collections.emptyList());

	handler.handleRequest(new HandlerRequest<>(request, EmptyMessageParameters.getInstance(), Collections.emptyMap(), Collections.emptyMap(), Collections.singleton(jobGraphFile.toFile())), mockGateway)
		.get();
}
 

@Test
public void testClusterShutdownWhenApplicationSucceeds() throws Exception {
	// we're "listening" on this to be completed to verify that the cluster
	// is being shut down from the ApplicationDispatcherBootstrap
	final CompletableFuture<ApplicationStatus> externalShutdownFuture = new CompletableFuture<>();

	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
			.setRequestJobStatusFunction(jobId -> CompletableFuture.completedFuture(JobStatus.FINISHED))
			.setRequestJobResultFunction(jobId -> CompletableFuture.completedFuture(createSuccessfulJobResult(jobId)))
			.setClusterShutdownFunction((status) -> {
				externalShutdownFuture.complete(status);
				return CompletableFuture.completedFuture(Acknowledge.get());
			});

	ApplicationDispatcherBootstrap bootstrap = createApplicationDispatcherBootstrap(3);

	final CompletableFuture<Acknowledge> shutdownFuture =
			bootstrap.runApplicationAndShutdownClusterAsync(dispatcherBuilder.build(), scheduledExecutor);

	// wait until the bootstrap "thinks" it's done
	shutdownFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);

	// verify that the dispatcher is actually being shut down
	assertThat(externalShutdownFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS), is(ApplicationStatus.SUCCEEDED));
}
 

@Test
public void testClusterShutdownWhenSubmissionFails() throws Exception {
	// we're "listening" on this to be completed to verify that the cluster
	// is being shut down from the ApplicationDispatcherBootstrap
	final CompletableFuture<ApplicationStatus> externalShutdownFuture = new CompletableFuture<>();

	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> {
				throw new FlinkRuntimeException("Nope!");
			})
			.setClusterShutdownFunction((status) -> {
				externalShutdownFuture.complete(status);
				return CompletableFuture.completedFuture(Acknowledge.get());
			});

	ApplicationDispatcherBootstrap bootstrap = createApplicationDispatcherBootstrap(3);

	final CompletableFuture<Acknowledge> shutdownFuture =
			bootstrap.runApplicationAndShutdownClusterAsync(dispatcherBuilder.build(), scheduledExecutor);

	// wait until the bootstrap "thinks" it's done
	shutdownFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);

	// verify that the dispatcher is actually being shut down
	assertThat(externalShutdownFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS), is(ApplicationStatus.FAILED));
}
 

@Test
public void testJobIdDefaultsToZeroWithHa() throws Throwable {
	final Configuration configurationUnderTest = getConfiguration();
	configurationUnderTest.set(HighAvailabilityOptions.HA_MODE, HighAvailabilityMode.ZOOKEEPER.name());

	final CompletableFuture<JobID> submittedJobId = new CompletableFuture<>();

	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> {
				submittedJobId.complete(jobGraph.getJobID());
				return CompletableFuture.completedFuture(Acknowledge.get());
			})
			.setRequestJobStatusFunction(jobId -> CompletableFuture.completedFuture(JobStatus.FINISHED))
			.setRequestJobResultFunction(jobId -> CompletableFuture.completedFuture(createSuccessfulJobResult(jobId)));

	final CompletableFuture<Void> applicationFuture =
			runApplication(dispatcherBuilder, configurationUnderTest, 1);

	applicationFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);

	assertThat(submittedJobId.get(TIMEOUT_SECONDS, TimeUnit.SECONDS), is(new JobID(0L, 0L)));
}
 

@Test
public void testStaticJobId() throws Throwable {
	final JobID testJobID = new JobID(0, 2);

	final Configuration configurationUnderTest = getConfiguration();
	configurationUnderTest.set(PipelineOptionsInternal.PIPELINE_FIXED_JOB_ID, testJobID.toHexString());

	final CompletableFuture<JobID> submittedJobId = new CompletableFuture<>();

	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> {
				submittedJobId.complete(jobGraph.getJobID());
				return CompletableFuture.completedFuture(Acknowledge.get());
			})
			.setRequestJobStatusFunction(jobId -> CompletableFuture.completedFuture(JobStatus.FINISHED))
			.setRequestJobResultFunction(jobId -> CompletableFuture.completedFuture(createSuccessfulJobResult(jobId)));

	final CompletableFuture<Void> applicationFuture =
			runApplication(dispatcherBuilder, configurationUnderTest, 1);

	applicationFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);

	assertThat(submittedJobId.get(TIMEOUT_SECONDS, TimeUnit.SECONDS), is(new JobID(0L, 2L)));
}
 

@Test
public void testStaticJobIdWithHa() throws Throwable {
	final JobID testJobID = new JobID(0, 2);

	final Configuration configurationUnderTest = getConfiguration();
	configurationUnderTest.set(PipelineOptionsInternal.PIPELINE_FIXED_JOB_ID, testJobID.toHexString());
	configurationUnderTest.set(HighAvailabilityOptions.HA_MODE, HighAvailabilityMode.ZOOKEEPER.name());

	final CompletableFuture<JobID> submittedJobId = new CompletableFuture<>();

	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> {
				submittedJobId.complete(jobGraph.getJobID());
				return CompletableFuture.completedFuture(Acknowledge.get());
			})
			.setRequestJobStatusFunction(jobId -> CompletableFuture.completedFuture(JobStatus.FINISHED))
			.setRequestJobResultFunction(jobId -> CompletableFuture.completedFuture(createSuccessfulJobResult(jobId)));

	final CompletableFuture<Void> applicationFuture =
			runApplication(dispatcherBuilder, configurationUnderTest, 1);

	applicationFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);

	assertThat(submittedJobId.get(TIMEOUT_SECONDS, TimeUnit.SECONDS), is(new JobID(0L, 2L)));
}
 

@Test
public void testApplicationIsStoppedWhenStoppingBootstrap() throws Exception {
	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
			.setRequestJobStatusFunction(jobId -> CompletableFuture.completedFuture(JobStatus.RUNNING));

	ApplicationDispatcherBootstrap bootstrap = createApplicationDispatcherBootstrap(3);

	final CompletableFuture<Acknowledge> shutdownFuture =
			bootstrap.runApplicationAndShutdownClusterAsync(dispatcherBuilder.build(), scheduledExecutor);

	ScheduledFuture<?> applicationExecutionFuture = bootstrap.getApplicationExecutionFuture();

	bootstrap.stop();

	// wait until the bootstrap "thinks" it's done
	shutdownFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);

	// verify that the application task is being cancelled
	assertThat(applicationExecutionFuture.isCancelled(), is(true));
}
 

@Test
public void testDispatcherIsCancelledWhenOneJobIsCancelled() throws Exception {
	final CompletableFuture<ApplicationStatus> clusterShutdownStatus = new CompletableFuture<>();

	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
			.setRequestJobStatusFunction(jobId -> CompletableFuture.completedFuture(JobStatus.CANCELED))
			.setClusterShutdownFunction((status) -> {
				clusterShutdownStatus.complete(status);
				return CompletableFuture.completedFuture(Acknowledge.get());
			})
			.setRequestJobResultFunction(jobId -> CompletableFuture.completedFuture(createCancelledJobResult(jobId)));

	ApplicationDispatcherBootstrap bootstrap = createApplicationDispatcherBootstrap(3);

	final CompletableFuture<Acknowledge> shutdownFuture =
			bootstrap.runApplicationAndShutdownClusterAsync(dispatcherBuilder.build(), scheduledExecutor);

	// wait until the bootstrap "thinks" it's done, also makes sure that we don't
	// fail the future exceptionally with a JobCancelledException
	shutdownFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);

	assertThat(clusterShutdownStatus.get(TIMEOUT_SECONDS, TimeUnit.SECONDS), is(ApplicationStatus.CANCELED));
}
 

@Test
public void testRejectionOnCountMismatch() throws Exception {
	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();
	try (ObjectOutputStream objectOut = new ObjectOutputStream(Files.newOutputStream(jobGraphFile))) {
		objectOut.writeObject(new JobGraph("testjob"));
	}
	final Path countExceedingFile = TEMPORARY_FOLDER.newFile().toPath();

	TestingDispatcherGateway.Builder builder = new TestingDispatcherGateway.Builder();
	builder
		.setBlobServerPort(blobServer.getPort())
		.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
		.setHostname("localhost");
	DispatcherGateway mockGateway = builder.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.getFileName().toString(), Collections.emptyList(), Collections.emptyList());

	try {
		handler.handleRequest(new HandlerRequest<>(request, EmptyMessageParameters.getInstance(), Collections.emptyMap(), Collections.emptyMap(), Arrays.asList(jobGraphFile.toFile(), countExceedingFile.toFile())), mockGateway)
			.get();
	} catch (Exception e) {
		ExceptionUtils.findThrowable(e, candidate -> candidate instanceof RestHandlerException && candidate.getMessage().contains("count"));
	}
}
 

@Test
public void testClusterShutdownWhenStoppingBootstrap() throws Exception {
	// we're "listening" on this to be completed to verify that the cluster
	// is being shut down from the ApplicationDispatcherBootstrap
	final CompletableFuture<ApplicationStatus> externalShutdownFuture = new CompletableFuture<>();

	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
			.setRequestJobStatusFunction(jobId -> CompletableFuture.completedFuture(JobStatus.RUNNING))
			.setClusterShutdownFunction((status) -> {
				externalShutdownFuture.complete(status);
				return CompletableFuture.completedFuture(Acknowledge.get());
			});

	ApplicationDispatcherBootstrap bootstrap = createApplicationDispatcherBootstrap(2);

	final CompletableFuture<Acknowledge> shutdownFuture =
			bootstrap.runApplicationAndShutdownClusterAsync(dispatcherBuilder.build(), scheduledExecutor);

	bootstrap.stop();

	// wait until the bootstrap "thinks" it's done
	shutdownFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);

	// verify that the dispatcher is actually being shut down
	assertThat(externalShutdownFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS), is(ApplicationStatus.UNKNOWN));
}
 

private CompletableFuture<Void> runApplication(
		TestingDispatcherGateway.Builder dispatcherBuilder,
		Configuration configuration,
		int noOfJobs) throws FlinkException {

	final PackagedProgram program = getProgram(noOfJobs);

	final ApplicationDispatcherBootstrap bootstrap =
			new ApplicationDispatcherBootstrap(
					program, Collections.emptyList(), configuration);

	return bootstrap.fixJobIdAndRunApplicationAsync(
			dispatcherBuilder.build(),
			scheduledExecutor);
}
 

@Test
public void testApplicationSucceedsWhenAllJobsSucceed() throws Exception {
	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
			.setRequestJobStatusFunction(jobId -> CompletableFuture.completedFuture(JobStatus.FINISHED))
			.setRequestJobResultFunction(jobId -> CompletableFuture.completedFuture(createSuccessfulJobResult(jobId)));

	final CompletableFuture<Void> applicationFuture = runApplication(dispatcherBuilder, 3);

	// this would block indefinitely if the applications don't finish
	applicationFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
}
 

@Test
public void testApplicationFailsAsSoonAsOneJobFails() throws Throwable {
	final ConcurrentLinkedDeque<JobID> submittedJobIds = new ConcurrentLinkedDeque<>();

	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> {
				submittedJobIds.add(jobGraph.getJobID());
				return CompletableFuture.completedFuture(Acknowledge.get());
			})
			.setRequestJobStatusFunction(jobId -> {
				// we only fail one of the jobs, the first one, the others will "keep" running
				// indefinitely
				if (jobId.equals(submittedJobIds.peek())) {
					return CompletableFuture.completedFuture(JobStatus.FAILED);
				}
				// never finish the other jobs
				return CompletableFuture.completedFuture(JobStatus.RUNNING);
			})
			.setRequestJobResultFunction(jobId -> {
				// we only fail one of the jobs, the first one, the other will "keep" running
				// indefinitely. If we didn't have this the test would hang forever.
				if (jobId.equals(submittedJobIds.peek())) {
					return CompletableFuture.completedFuture(createFailedJobResult(jobId));
				}
				// never finish the other jobs
				return new CompletableFuture<>();
			});

	final CompletableFuture<Void> applicationFuture = runApplication(dispatcherBuilder, 2);

	assertException(applicationFuture, JobExecutionException.class);
}
 

@Test
public void testExceptionThrownWhenApplicationContainsNoJobs() throws Throwable {
	final TestingDispatcherGateway.Builder dispatcherBuilder = new TestingDispatcherGateway.Builder()
			.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()));

	final CompletableFuture<Void> applicationFuture =
			runApplication(dispatcherBuilder, 0);

	assertException(applicationFuture, ApplicationExecutionException.class);
}
 

@Test
public void testRejectionOnCountMismatch() throws Exception {
	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();
	try (ObjectOutputStream objectOut = new ObjectOutputStream(Files.newOutputStream(jobGraphFile))) {
		objectOut.writeObject(new JobGraph("testjob"));
	}
	final Path countExceedingFile = TEMPORARY_FOLDER.newFile().toPath();

	TestingDispatcherGateway.Builder builder = new TestingDispatcherGateway.Builder();
	builder
		.setBlobServerPort(blobServer.getPort())
		.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
		.setHostname("localhost");
	DispatcherGateway mockGateway = builder.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.getFileName().toString(), Collections.emptyList(), Collections.emptyList());

	try {
		handler.handleRequest(new HandlerRequest<>(request, EmptyMessageParameters.getInstance(), Collections.emptyMap(), Collections.emptyMap(), Arrays.asList(jobGraphFile.toFile(), countExceedingFile.toFile())), mockGateway)
			.get();
	} catch (Exception e) {
		ExceptionUtils.findThrowable(e, candidate -> candidate instanceof RestHandlerException && candidate.getMessage().contains("count"));
	}
}
 

@Test
public void testRejectionOnCountMismatch() throws Exception {
	final Path jobGraphFile = TEMPORARY_FOLDER.newFile().toPath();
	try (ObjectOutputStream objectOut = new ObjectOutputStream(Files.newOutputStream(jobGraphFile))) {
		objectOut.writeObject(new JobGraph("testjob"));
	}
	final Path countExceedingFile = TEMPORARY_FOLDER.newFile().toPath();

	TestingDispatcherGateway.Builder builder = new TestingDispatcherGateway.Builder();
	builder
		.setBlobServerPort(blobServer.getPort())
		.setSubmitFunction(jobGraph -> CompletableFuture.completedFuture(Acknowledge.get()))
		.setHostname("localhost");
	DispatcherGateway mockGateway = builder.build();

	JobSubmitHandler handler = new JobSubmitHandler(
		() -> CompletableFuture.completedFuture(mockGateway),
		RpcUtils.INF_TIMEOUT,
		Collections.emptyMap(),
		TestingUtils.defaultExecutor(),
		configuration);

	JobSubmitRequestBody request = new JobSubmitRequestBody(jobGraphFile.getFileName().toString(), Collections.emptyList(), Collections.emptyList());

	try {
		handler.handleRequest(new HandlerRequest<>(request, EmptyMessageParameters.getInstance(), Collections.emptyMap(), Collections.emptyMap(), Arrays.asList(jobGraphFile.toFile(), countExceedingFile.toFile())), mockGateway)
			.get();
	} catch (Exception e) {
		ExceptionUtils.findThrowable(e, candidate -> candidate instanceof RestHandlerException && candidate.getMessage().contains("count"));
	}
}
 

private CompletableFuture<Void> runApplication(
		TestingDispatcherGateway.Builder dispatcherBuilder,
		int noOfJobs) throws FlinkException {

	return runApplication(dispatcherBuilder, getConfiguration(), noOfJobs);
}