@Test
public void testFindTasksWithLegacyMigration() {

    long deadline1 = titusRuntime.getClock().wallTime() + 1_000;
    long legacyDeadline2 = titusRuntime.getClock().wallTime() + 2_000;
    Task legacyTask = toLegacyTask(TASK2, legacyDeadline2);
    TaskQueryResult queryResult = TaskQueryResult.newBuilder()
            .addItems(TASK1)
            .addItems(legacyTask)
            .build();
    List<TaskRelocationPlan> relocationPlans = Collections.singletonList(newRelocationPlan(TASK1, deadline1));
    when(relocationServiceClient.findTaskRelocationPlans(asSet(TASK1.getId(), TASK2.getId()))).thenReturn(Mono.just(relocationPlans));
    TaskQueryResult merged = taskRelocationDataInjector.injectIntoTaskQueryResult(Observable.just(queryResult)).toBlocking().first();
    assertThat(merged.getItems(0).getMigrationDetails().getNeedsMigration()).isTrue();
    assertThat(merged.getItems(0).getMigrationDetails().getDeadline()).isEqualTo(deadline1);
    assertThat(merged.getItems(1)).isEqualTo(legacyTask);
}