private static ClusterClient<String> createClusterClient() throws Exception { final ClusterClient<String> clusterClient = mock(ClusterClient.class); when(clusterClient.listJobs()).thenReturn(CompletableFuture.completedFuture(Arrays.asList( new JobStatusMessage(new JobID(), "job1", JobStatus.RUNNING, 1L), new JobStatusMessage(new JobID(), "job2", JobStatus.CREATED, 1L), new JobStatusMessage(new JobID(), "job3", JobStatus.SUSPENDING, 3L), new JobStatusMessage(new JobID(), "job4", JobStatus.SUSPENDING, 2L), new JobStatusMessage(new JobID(), "job5", JobStatus.FINISHED, 3L) ))); return clusterClient; } }
@Test public void testStop() throws Exception { // test stop properly JobID jid = new JobID(); String jidString = jid.toString(); String[] parameters = { jidString }; final ClusterClient<String> clusterClient = createClusterClient(null); MockedCliFrontend testFrontend = new MockedCliFrontend(clusterClient); testFrontend.stop(parameters); Mockito.verify(clusterClient, times(1)).stop(any(JobID.class)); }
/** * Creates a new Execution Environment. */ protected ExecutionEnvironment() { jobID = JobID.generate(); }
/** * Tests cancelling with the savepoint option. */ @Test public void testCancelWithSavepoint() throws Exception { { // Cancel with savepoint (no target directory) JobID jid = new JobID(); String[] parameters = { "-s", jid.toString() }; final ClusterClient<String> clusterClient = createClusterClient(); MockedCliFrontend testFrontend = new MockedCliFrontend(clusterClient); testFrontend.cancel(parameters); Mockito.verify(clusterClient, times(1)) .cancelWithSavepoint(any(JobID.class), isNull(String.class)); } { // Cancel with savepoint (with target directory) JobID jid = new JobID(); String[] parameters = { "-s", "targetDirectory", jid.toString() }; final ClusterClient<String> clusterClient = createClusterClient(); MockedCliFrontend testFrontend = new MockedCliFrontend(clusterClient); testFrontend.cancel(parameters); Mockito.verify(clusterClient, times(1)) .cancelWithSavepoint(any(JobID.class), notNull(String.class)); } }
StreamTask<?, ?> streamTask = spy(new EmptyStreamTask(mockEnvironment)); CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, timestamp); StreamOperator<?> streamOperator1 = mock(StreamOperator.class); StreamOperator<?> streamOperator2 = mock(StreamOperator.class); StreamOperator<?> streamOperator3 = mock(StreamOperator.class); RunnableFuture<SnapshotResult<OperatorStateHandle>> failingFuture = mock(RunnableFuture.class); when(failingFuture.get()).thenThrow(new ExecutionException(new Exception("Test exception"))); when(operatorSnapshotResult3.getOperatorStateRawFuture()).thenReturn(failingFuture); when(streamOperator1.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenReturn(operatorSnapshotResult1); when(streamOperator2.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenReturn(operatorSnapshotResult2); when(streamOperator3.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenReturn(operatorSnapshotResult3); Whitebox.setInternalState(streamTask, "cancelables", new CloseableRegistry()); Whitebox.setInternalState(streamTask, "asyncOperationsThreadPool", newDirectExecutorService()); Whitebox.setInternalState(streamTask, "configuration", new StreamConfig(new Configuration())); Whitebox.setInternalState(streamTask, "checkpointStorage", new MemoryBackendCheckpointStorage(new JobID(), null, null, Integer.MAX_VALUE));
final OneShotLatch completeSubtask = new OneShotLatch(); Environment mockEnvironment = spy(new MockEnvironmentBuilder().build()); thenAnswer((Answer<OperatorSnapshotFinalizer>) invocation -> { createSubtask.trigger(); completeSubtask.await(); CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, timestamp); final StreamOperator<?> streamOperator = mock(StreamOperator.class); final OperatorID operatorID = new OperatorID(); when(streamOperator.getOperatorID()).thenReturn(operatorID); KeyedStateHandle managedKeyedStateHandle = mock(KeyedStateHandle.class); DoneFuture.of(SnapshotResult.of(rawOperatorStateHandle))); when(streamOperator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenReturn(operatorSnapshotResult); when(operatorChain.getAllOperators()).thenReturn(streamOperators); CheckpointStorage checkpointStorage = new MemoryBackendCheckpointStorage(new JobID(), null, null, Integer.MAX_VALUE); Whitebox.setInternalState(streamTask, "cancelables", new CloseableRegistry()); Whitebox.setInternalState(streamTask, "asyncOperationsThreadPool", executor); Whitebox.setInternalState(streamTask, "configuration", new StreamConfig(new Configuration())); Whitebox.setInternalState(streamTask, "checkpointStorage", checkpointStorage);
final OneShotLatch completeAcknowledge = new OneShotLatch(); CheckpointResponder checkpointResponder = mock(CheckpointResponder.class); doAnswer(new Answer() { @Override public Object answer(InvocationOnMock invocation) throws Throwable { new JobID(1L, 2L), new ExecutionAttemptID(1L, 2L), mock(TaskLocalStateStoreImpl.class), null, checkpointResponder); StreamOperator<?> streamOperator = mock(StreamOperator.class); when(streamOperator.getOperatorID()).thenReturn(new OperatorID(42, 42)); KeyedStateHandle managedKeyedStateHandle = mock(KeyedStateHandle.class); DoneFuture.of(SnapshotResult.of(rawOperatorStateHandle))); when(streamOperator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenReturn(operatorSnapshotResult); OperatorChain<Void, AbstractStreamOperator<Void>> operatorChain = mock(OperatorChain.class); when(operatorChain.getAllOperators()).thenReturn(streamOperators); CheckpointStorage checkpointStorage = new MemoryBackendCheckpointStorage(new JobID(), null, null, Integer.MAX_VALUE);
new BlobCacheService(mock(PermanentBlobCache.class), mock(TransientBlobCache.class)); LibraryCacheManager libCache = mock(LibraryCacheManager.class); when(libCache.getClassLoader(any(JobID.class))).thenReturn(StreamTaskTest.class.getClassLoader()); ResultPartitionManager partitionManager = mock(ResultPartitionManager.class); NetworkEnvironment network = mock(NetworkEnvironment.class); when(network.getResultPartitionManager()).thenReturn(partitionManager); when(network.getDefaultIOMode()).thenReturn(IOManager.IOMode.SYNC); when(network.createKvStateTaskRegistry(any(JobID.class), any(JobVertexID.class))) .thenReturn(mock(TaskKvStateRegistry.class)); when(network.getTaskEventDispatcher()).thenReturn(taskEventDispatcher); new JobID(), "Job Name", new SerializedValue<>(new ExecutionConfig()), new Configuration(), Collections.emptyList(), Collections.emptyList());
final long timestamp = 1L; Environment mockEnvironment = spy(new MockEnvironmentBuilder().build()); final List<SubtaskState> checkpointResult = new ArrayList<>(1); new JobID(1L, 2L), new ExecutionAttemptID(1L, 2L), mock(TaskLocalStateStoreImpl.class), checkpointResponder); when(mockEnvironment.getTaskStateManager()).thenReturn(taskStateManager); when(statelessOperator.getOperatorID()).thenReturn(operatorID); when(statelessOperator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))) .thenReturn(statelessOperatorSnapshotResult); Whitebox.setInternalState(streamTask, "operatorChain", operatorChain); Whitebox.setInternalState(streamTask, "cancelables", new CloseableRegistry()); Whitebox.setInternalState(streamTask, "configuration", new StreamConfig(new Configuration())); Whitebox.setInternalState(streamTask, "asyncOperationsThreadPool", Executors.newCachedThreadPool()); Whitebox.setInternalState(streamTask, "checkpointStorage", new MemoryBackendCheckpointStorage(new JobID(), null, null, Integer.MAX_VALUE));
final Configuration taskConfiguration = new Configuration(); final StreamConfig streamConfig = new StreamConfig(taskConfiguration); final NoOpStreamOperator<Long> noOpStreamOperator = new NoOpStreamOperator<>(); new JobID(), "Test Job", new SerializedValue<>(new ExecutionConfig()), new Configuration(), Collections.emptyList(), Collections.emptyList()); final NetworkEnvironment networkEnv = mock(NetworkEnvironment.class); when(networkEnv.createKvStateTaskRegistry(any(JobID.class), any(JobVertexID.class))).thenReturn(mock(TaskKvStateRegistry.class)); when(networkEnv.getTaskEventDispatcher()).thenReturn(taskEventDispatcher); Executors.directExecutor()); CompletableFuture<Void> taskRun = CompletableFuture.runAsync( () -> task.run(), TestingUtils.defaultExecutor()); taskRun.get();
NetworkEnvironment networkEnvironment = mock(NetworkEnvironment.class); when(networkEnvironment.createKvStateTaskRegistry(any(JobID.class), any(JobVertexID.class))) .thenReturn(mock(TaskKvStateRegistry.class)); when(networkEnvironment.getTaskEventDispatcher()).thenReturn(taskEventDispatcher); new JobID(), "test job name", new SerializedValue<>(new ExecutionConfig()), new Configuration(), Collections.emptyList(), Collections.emptyList());
new SlotRequest(new JobID(), new AllocationID(), resourceProfile1, taskHost)); return null; }); registerSlotRequestFuture.get(); doReturn(Collections.singletonList(Collections.singletonList(resourceManager.getContainerRequest()))) .when(mockResourceManagerClient).getMatchingRequests(any(Priority.class), anyString(), any(Resource.class)); verify(mockResourceManagerClient).addContainerRequest(any(AMRMClient.ContainerRequest.class)); verify(mockNMClient).startContainer(eq(testingContainer), any(ContainerLaunchContext.class)); hardwareDescription, Time.seconds(10L)) .thenCompose( (RegistrationResponse response) -> { assertThat(response, instanceOf(TaskExecutorRegistrationSuccess.class)); Time.seconds(10L)); }) .handleAsync( (Acknowledge ignored, Throwable throwable) -> rmServices.slotManager.getNumberRegisteredSlots(), resourceManager.getMainThreadExecutorForTesting());
CompletableFuture<?> registerSlotRequestFuture = resourceManager.runInMainThread(() -> { rmServices.slotManager.registerSlotRequest( new SlotRequest(new JobID(), new AllocationID(), resourceProfile1, taskHost)); return null; }); registerSlotRequestFuture.get(); doReturn(Collections.singletonList(Collections.singletonList(resourceManager.getContainerRequest()))) .when(mockResourceManagerClient).getMatchingRequests(any(Priority.class), anyString(), any(Resource.class)); verify(mockResourceManagerClient).addContainerRequest(any(AMRMClient.ContainerRequest.class)); verify(mockResourceManagerClient).removeContainerRequest(any(AMRMClient.ContainerRequest.class)); verify(mockNMClient).startContainer(eq(testingContainer), any(ContainerLaunchContext.class));
@Test public void testDeployerWithIsolatedConfiguration() throws Exception { YarnClusterConfiguration clusterConf = mock(YarnClusterConfiguration.class); doReturn(new YarnConfiguration()).when(clusterConf).conf(); ScheduledExecutorService executor = mock(ScheduledExecutorService.class); Configuration flinkConf = new Configuration(); YarnClient client = mock(YarnClient.class); JobDeployer deploy = new JobDeployer(clusterConf, client, executor, flinkConf); AthenaXYarnClusterDescriptor desc = mock(AthenaXYarnClusterDescriptor.class); YarnClusterClient clusterClient = mock(YarnClusterClient.class); doReturn(clusterClient).when(desc).deploy(); ActorGateway actorGateway = mock(ActorGateway.class); doReturn(actorGateway).when(clusterClient).getJobManagerGateway(); doReturn(Future$.MODULE$.successful(null)).when(actorGateway).ask(any(), any()); JobGraph jobGraph = mock(JobGraph.class); doReturn(JobID.generate()).when(jobGraph).getJobID(); deploy.start(desc, jobGraph); verify(clusterClient).runDetached(jobGraph, null); } }
@Test public void testTriggerSavepointSuccess() throws Exception { replaceStdOutAndStdErr(); JobID jobId = new JobID(); String savepointPath = "expectedSavepointPath"; final ClusterClient<String> clusterClient = createClusterClient(savepointPath); try { MockedCliFrontend frontend = new MockedCliFrontend(clusterClient); String[] parameters = { jobId.toString() }; frontend.savepoint(parameters); verify(clusterClient, times(1)) .triggerSavepoint(eq(jobId), isNull(String.class)); assertTrue(buffer.toString().contains(savepointPath)); } finally { clusterClient.shutdown(); restoreStdOutAndStdErr(); } }
@Test public void testNonExistingJobRetrieval() throws Exception { final JobID jobID = new JobID(); try { client.requestJobResult(jobID).get(); fail(); } catch (Exception exception) { Optional<Throwable> expectedCause = ExceptionUtils.findThrowable(exception, candidate -> candidate.getMessage() != null && candidate.getMessage().contains("Could not find Flink job")); if (!expectedCause.isPresent()) { throw exception; } } }
@Test public void testClusterClientSavepoint() throws Exception { Configuration config = new Configuration(); config.setString(JobManagerOptions.ADDRESS, "localhost"); JobID jobID = new JobID(); String savepointDirectory = "/test/directory"; String savepointPath = "/test/path"; TestSavepointActorGateway gateway = new TestSavepointActorGateway(jobID, savepointDirectory, savepointPath); TestClusterClient clusterClient = new TestClusterClient(config, gateway); try { CompletableFuture<String> pathFuture = clusterClient.triggerSavepoint(jobID, savepointDirectory); Assert.assertTrue(gateway.messageArrived); Assert.assertEquals(savepointPath, pathFuture.get()); } finally { clusterClient.shutdown(); } }
private static MapState<EmailId, EmailInformation> getMapState( String jobId, QueryableStateClient client, MapStateDescriptor<EmailId, EmailInformation> stateDescriptor) throws InterruptedException, ExecutionException { CompletableFuture<MapState<EmailId, EmailInformation>> resultFuture = client.getKvState( JobID.fromHexString(jobId), QsConstants.QUERY_NAME, QsConstants.KEY, // which key of the keyed state to access BasicTypeInfo.STRING_TYPE_INFO, stateDescriptor); return resultFuture.get(); } }
final JobID jobId = new JobID(); final String jobName = "Semi-Rebalance Test Job"; final Configuration cfg = new Configuration(); new NoRestartStrategy(), new RestartAllStrategy.Factory(), new TestingSlotProvider(ignored -> new CompletableFuture<>()), ExecutionGraph.class.getClassLoader(), VoidBlobWriter.getInstance(),