MustBeRelocatedSelfManagedTaskCollectorStep { public Map<String, TaskRelocationPlan> collectTasksThatMustBeRelocated() { Stopwatch stopwatch = Stopwatch.createStarted(); try { Map<String, TaskRelocationPlan> result = buildRelocationPlans(); metrics.onSuccess(result.size(), stopwatch.elapsed(TimeUnit.MILLISECONDS)); logger.debug("Step results: {}", result); return result; } catch (Exception e) { logger.error("Step processing error", e); metrics.onError(stopwatch.elapsed(TimeUnit.MILLISECONDS)); throw e; } } MustBeRelocatedSelfManagedTaskCollectorStep(NodeDataResolver nodeDataResolver,
                                                       ReadOnlyJobOperations jobOperations,
                                                       TitusRuntime titusRuntime); Map<String, TaskRelocationPlan> collectTasksThatMustBeRelocated();  }
@Test public void testMigrationOfTasksWithPolicy() { Job<BatchJobExt> job = TestDataFactory.newBatchJob("job1", 1, newSelfManagedDisruptionBudget(1_000)); relocationConnectorStubs.addJob(job); relocationConnectorStubs.place(TestDataFactory.REMOVABLE_INSTANCE_GROUP, jobOperations.getTasks().get(0)); Map<String, TaskRelocationPlan> result = step.collectTasksThatMustBeRelocated(); assertThat(result).hasSize(1); }
@Test public void testTaskWithNoDisruptionBudgetIsNotMigrated() { Job<BatchJobExt> job = TestDataFactory.newBatchJob("job1", 1, JobFunctions.getNoDisruptionBudgetMarker()); relocationConnectorStubs.addJob(job); relocationConnectorStubs.place(TestDataFactory.REMOVABLE_INSTANCE_GROUP, jobOperations.getTasks().get(0)); Map<String, TaskRelocationPlan> result = step.collectTasksThatMustBeRelocated(); assertThat(result).isEmpty(); }
@Test public void testOriginalPlanIsReturnedOnEachInvocation() { Job<BatchJobExt> job = TestDataFactory.newBatchJob("job1", 1, newSelfManagedDisruptionBudget(1_000)); relocationConnectorStubs.addJob(job); Task task = jobOperations.getTasks().get(0); relocationConnectorStubs.place(TestDataFactory.REMOVABLE_INSTANCE_GROUP, task); Map<String, TaskRelocationPlan> firstResult = step.collectTasksThatMustBeRelocated(); assertThat(firstResult).hasSize(1); TaskRelocationPlan first = firstResult.get(task.getId()); ((TestClock) titusRuntime.getClock()).advanceTime(1, TimeUnit.SECONDS); Map<String, TaskRelocationPlan> secondResult = step.collectTasksThatMustBeRelocated(); assertThat(secondResult).hasSize(1); TaskRelocationPlan second = secondResult.get(task.getId()); assertThat(first).isEqualTo(second); }