TaskTerminationExecutor { public Mono<Void> terminateTask(String taskId, String reason, String callerId) { return findAndVerifyJobAndTask(taskId, reason, callerId) .flatMap(jobTaskPair -> { Job<?> job = jobTaskPair.getLeft(); Task task = jobTaskPair.getRight(); return serializedInvoker .submit(doTerminateTask(taskId, reason, job, task, callerId)) .doOnSuccess(next -> onSuccessfulTermination(job, taskId, reason, callerId)) .doOnError(error -> onTerminationError(job, taskId, reason, callerId, error)); }); } TaskTerminationExecutor(EvictionServiceConfiguration configuration,
                            V3JobOperations jobOperations,
                            TitusQuotasManager quotasManager,
                            TitusRuntime titusRuntime,
                            Scheduler scheduler); Mono<Void> terminateTask(String taskId, String reason, String callerId);  }
@Test public void testTerminateTaskWithEnoughQuota() { Pair<Job, List<Task>> jobAndTask = createAndStartJobWithTasks(1); Job<?> job = jobAndTask.getLeft(); Task task = jobAndTask.getRight().get(0); when(quotasManager.findEvictionQuota(Reference.job(job.getId()))).thenReturn(Optional.of(EvictionQuota.jobQuota(job.getId(), 1, "Test"))); when(quotasManager.tryConsumeQuota(job, task)).thenReturn(ConsumptionResult.approved()); StepVerifier.withVirtualTime( () -> newTerminationExecutor().terminateTask(task.getId(), EVICTION_REASON, "callerContext") ).verifyComplete(); expectTaskTerminationEvent(task, true); verify(quotasManager, times(1)).tryConsumeQuota(job, task); }
@Test public void testTerminateTaskWithNoQuota() { Pair<Job, List<Task>> jobAndTask = createAndStartJobWithTasks(1); Job<?> job = jobAndTask.getLeft(); Task task = jobAndTask.getRight().get(0); when(quotasManager.tryConsumeQuota(job, task)).thenReturn(ConsumptionResult.rejected("no quota")); StepVerifier .withVirtualTime(() -> newTerminationExecutor().terminateTask(task.getId(), EVICTION_REASON, "callerContext")) .consumeErrorWith(e -> expectEvictionError(e, ErrorCode.NoQuota)) .verify(); expectTaskTerminationEvent(task, false); verify(quotasManager, times(1)).tryConsumeQuota(job, task); }