@Override
    public Optional<InputSplit> getNextInputSplit(String host, int attemptNumber) {
        final int index = nextInputSplitIndexToConsumeByAttempts.getOrDefault(attemptNumber, 0);
        checkState(index <= inputSplits.size());

        if (index < inputSplits.size()) {
            nextInputSplitIndexToConsumeByAttempts.put(attemptNumber, index + 1);
            return Optional.of(inputSplits.get(index));
        } else {
            final Optional<InputSplit> split = super.getNextInputSplit(host, attemptNumber);
            if (split.isPresent()) {
                nextInputSplitIndexToConsumeByAttempts.put(attemptNumber, index + 1);
            }
            return split;
        }
    }
@Test
    void testGetNextInputSplit() throws Exception {
        final TestInputSource source = new TestInputSource();
        final JobVertex jobVertex = ExecutionGraphTestUtils.createNoOpVertex(1);
        jobVertex.setInputSplitSource(source);

        final SpeculativeExecutionVertex ev = createSpeculativeExecutionVertex(jobVertex);

        final int numExecutions = 3;
        for (int i = 0; i < numExecutions - 1; ++i) {
            ev.createNewSpeculativeExecution(0);
        }
        final List<Execution> executions = new ArrayList<>(ev.getCurrentExecutions());

        final Map<Integer, List<InputSplit>> splitsOfAttempts = new HashMap<>();
        final Random rand = new Random();
        while (executions.size() > 0) {
            final int index = rand.nextInt(executions.size());
            final Execution execution = executions.get(index);
            final Optional<InputSplit> split = execution.getNextInputSplit();
            if (split.isPresent()) {
                splitsOfAttempts
                        .computeIfAbsent(execution.getAttemptNumber(), k -> new ArrayList<>())
                        .add(split.get());
            } else {
                executions.remove(index);
            }
        }

        assertThat(splitsOfAttempts).hasSize(numExecutions);
        assertThat(splitsOfAttempts.get(0)).containsExactlyInAnyOrder(source.splits);
        assertThat(splitsOfAttempts.get(1)).isEqualTo(splitsOfAttempts.get(0));
        assertThat(splitsOfAttempts.get(2)).isEqualTo(splitsOfAttempts.get(0));
    }