@Override
    public CompletableFuture<Collection<TaskManagerLocation>> getPreferredLocations(
            final ExecutionVertexID executionVertexId,
            final Set<ExecutionVertexID> producersToIgnore) {

        checkNotNull(executionVertexId);
        checkNotNull(producersToIgnore);

        final Collection<TaskManagerLocation> preferredLocationsBasedOnState =
                getPreferredLocationsBasedOnState(executionVertexId);
        if (!preferredLocationsBasedOnState.isEmpty()) {
            return CompletableFuture.completedFuture(preferredLocationsBasedOnState);
        }

        return getPreferredLocationsBasedOnInputs(executionVertexId, producersToIgnore);
    }
@Test
    void testInputLocationsChoosesInputOfFewerLocations() {
        final TestingInputsLocationsRetriever.Builder locationRetrieverBuilder =
                new TestingInputsLocationsRetriever.Builder();

        final ExecutionVertexID consumerId = new ExecutionVertexID(new JobVertexID(), 0);

        int parallelism1 = 3;
        final JobVertexID jobVertexId1 = new JobVertexID();
        final List<ExecutionVertexID> producers1 = new ArrayList<>(parallelism1);
        for (int i = 0; i < parallelism1; i++) {
            final ExecutionVertexID producerId = new ExecutionVertexID(jobVertexId1, i);
            producers1.add(producerId);
        }
        locationRetrieverBuilder.connectConsumerToProducers(consumerId, producers1);

        final JobVertexID jobVertexId2 = new JobVertexID();
        int parallelism2 = 5;
        final List<ExecutionVertexID> producers2 = new ArrayList<>(parallelism2);
        for (int i = 0; i < parallelism2; i++) {
            final ExecutionVertexID producerId = new ExecutionVertexID(jobVertexId2, i);
            producers2.add(producerId);
        }
        locationRetrieverBuilder.connectConsumerToProducers(consumerId, producers2);

        final TestingInputsLocationsRetriever inputsLocationsRetriever =
                locationRetrieverBuilder.build();

        final List<TaskManagerLocation> expectedLocations = new ArrayList<>(parallelism1);
        for (int i = 0; i < parallelism1; i++) {
            inputsLocationsRetriever.assignTaskManagerLocation(producers1.get(i));
            expectedLocations.add(
                    inputsLocationsRetriever
                            .getTaskManagerLocation(producers1.get(i))
                            .get()
                            .getNow(null));
        }

        for (int i = 0; i < parallelism2; i++) {
            inputsLocationsRetriever.assignTaskManagerLocation(producers2.get(i));
        }

        final PreferredLocationsRetriever locationsRetriever =
                new DefaultPreferredLocationsRetriever(
                        id -> Optional.empty(), inputsLocationsRetriever);

        final CompletableFuture<Collection<TaskManagerLocation>> preferredLocations =
                locationsRetriever.getPreferredLocations(consumerId, Collections.emptySet());

        assertThat(preferredLocations.getNow(null))
                .containsExactlyInAnyOrderElementsOf(expectedLocations);
    }