@Override
    public Optional<FileSourceSplit> getNext(@Nullable String host) {
        // for a null host, we always return a remote split
        if (StringUtils.isNullOrWhitespaceOnly(host)) {
            final Optional<FileSourceSplit> split = getRemoteSplit();
            if (split.isPresent()) {
                LOG.info("Assigning split to non-localized request: {}", split);
            }
            return split;
        }

        host = normalizeHostName(host);

        // for any non-null host, we take the list of non-null splits
        final LocatableSplitChooser localSplits =
                localPerHost.computeIfAbsent(
                        host, (theHost) -> buildChooserForHost(theHost, unassigned));

        final SplitWithInfo localSplit =
                localSplits.getNextUnassignedMinLocalCountSplit(unassigned);
        if (localSplit != null) {
            checkState(
                    unassigned.remove(localSplit),
                    "Selected split has already been assigned. This should not happen!");
            LOG.info(
                    "Assigning local split to requesting host '{}': {}",
                    host,
                    localSplit.getSplit());
            localAssignments.inc();
            return Optional.of(localSplit.getSplit());
        }

        // we did not find a local split, return a remote split
        final Optional<FileSourceSplit> remoteSplit = getRemoteSplit();
        if (remoteSplit.isPresent()) {
            LOG.info("Assigning remote split to requesting host '{}': {}", host, remoteSplit);
        }
        return remoteSplit;
    }
@Test
    void testAssignmentOfManySplitsRandomly() {
        final long seed = Calendar.getInstance().getTimeInMillis();

        final int numSplits = 1000;
        final String[] splitHosts = new String[256];
        final String[] requestingHosts = new String[256];
        final Random rand = new Random(seed);

        for (int i = 0; i < splitHosts.length; i++) {
            splitHosts[i] = "localHost" + i;
        }
        for (int i = 0; i < requestingHosts.length; i++) {
            if (i % 2 == 0) {
                requestingHosts[i] = "localHost" + i;
            } else {
                requestingHosts[i] = "remoteHost" + i;
            }
        }

        String[] stringArray = {};
        Set<String> hosts = new HashSet<>();
        Set<FileSourceSplit> splits = new HashSet<>();
        for (int i = 0; i < numSplits; i++) {
            while (hosts.size() < 3) {
                hosts.add(splitHosts[rand.nextInt(splitHosts.length)]);
            }
            splits.add(createSplit(i, hosts.toArray(stringArray)));
            hosts.clear();
        }

        final LocalityAwareSplitAssigner ia = new LocalityAwareSplitAssigner(splits);

        for (int i = 0; i < numSplits; i++) {
            final Optional<FileSourceSplit> split =
                    ia.getNext(requestingHosts[rand.nextInt(requestingHosts.length)]);
            assertThat(split).isPresent();
            assertThat(splits.remove(split.get())).isTrue();
        }

        assertThat(splits).isEmpty();
        assertThat(ia.getNext("testHost")).isNotPresent();
    }