TaskRelocationDataInjector { Observable<Task> injectIntoTask(String taskId, Observable<Task> taskObservable) { if (!featureActivationConfiguration.isMergingTaskMigrationPlanInGatewayEnabled()) { return taskObservable; } if (shouldUseRelocationCache()) { return taskObservable.map(task -> newTaskWithRelocationPlan(task, relocationDataReplicator.getCurrent().getPlans().get(taskId))); } Observable<Optional<TaskRelocationPlan>> relocationPlanResolver = ReactorExt.toObservable(relocationServiceClient.findTaskRelocationPlan(taskId)) .timeout(getTaskRelocationTimeout(), TimeUnit.MILLISECONDS, scheduler) .doOnError(error -> logger.info("Could not resolve task relocation status for task: taskId={}, error={}", taskId, ExceptionExt.toMessageChain(error))) .onErrorReturn(e -> Optional.empty()); return Observable.zip( taskObservable, relocationPlanResolver, (task, planOpt) -> planOpt.map(plan -> newTaskWithRelocationPlan(task, plan)).orElse(task) ); } @Inject  TaskRelocationDataInjector(
            GrpcClientConfiguration configuration,
            JobManagerConfiguration jobManagerConfiguration,
            FeatureActivationConfiguration featureActivationConfiguration,
            RelocationServiceClient relocationServiceClient,
            RelocationDataReplicator relocationDataReplicator); @VisibleForTesting  TaskRelocationDataInjector(
            GrpcClientConfiguration configuration,
            JobManagerConfiguration jobManagerConfiguration,
            FeatureActivationConfiguration featureActivationConfiguration,
            RelocationServiceClient relocationServiceClient,
            RelocationDataReplicator relocationDataReplicator,
            Scheduler scheduler);   }
@Test public void testFindTaskWithRelocationDeadline() { long deadlineTimestamp = titusRuntime.getClock().wallTime() + 1_000; when(relocationServiceClient.findTaskRelocationPlan(TASK1.getId())).thenReturn( Mono.just(Optional.of(newRelocationPlan(TASK1, deadlineTimestamp))) ); Task merged = taskRelocationDataInjector.injectIntoTask(TASK1.getId(), Observable.just(TASK1)).toBlocking().first(); assertThat(merged.getMigrationDetails().getNeedsMigration()).isTrue(); assertThat(merged.getMigrationDetails().getDeadline()).isEqualTo(deadlineTimestamp); }
@Test public void testFindTaskWithoutRelocationDeadline() { when(relocationServiceClient.findTaskRelocationPlan(TASK1.getId())).thenReturn(Mono.just(Optional.empty())); Task merged = taskRelocationDataInjector.injectIntoTask(TASK1.getId(), Observable.just(TASK1)).toBlocking().first(); assertThat(merged.getMigrationDetails().getNeedsMigration()).isFalse(); }
@Test public void testFindTaskWithLegacyMigration() { long deadlineTimestamp = titusRuntime.getClock().wallTime() + 1_000; Task legacyTask = toLegacyTask(TASK1, deadlineTimestamp); when(relocationServiceClient.findTaskRelocationPlan(TASK1.getId())).thenReturn(Mono.just(Optional.empty())); taskRelocationDataInjector.injectIntoTask(legacyTask.getId(), Observable.just(legacyTask)).toBlocking().first(); assertThat(legacyTask).isEqualToComparingFieldByField(legacyTask); }
@Test public void testFindTaskWithRelocationDataFetchTimeout() { when(relocationServiceClient.findTaskRelocationPlan(TASK1.getId())).thenReturn(Mono.never()); ExtTestSubscriber<Task> testSubscriber = new ExtTestSubscriber<>(); taskRelocationDataInjector.injectIntoTask(TASK1.getId(), Observable.just(TASK1)).subscribe(testSubscriber); testSubscriber.assertOpen(); testScheduler.advanceTimeBy(RELOCATION_TIMEOUT_MS, TimeUnit.MILLISECONDS); assertThat(testSubscriber.takeNext()).isEqualTo(TASK1); testSubscriber.assertOnCompleted(); }