iterator (Iterator<T>) that produces globally sorted sequence from the given N iterators

SortedJoiningIterator.java

package com.exercise;

import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.PriorityQueue;

public class SortedJoiningIterator<T> implements Iterator<T> {
    private final PriorityQueue<Head<T>> iterators;
    private Iterator<T> lastIt;

    private static class Head<T> {
        private final Iterator<T> it;
        private T value;
        private boolean nextIsCalled = false;

        private Head(Iterator<T> it) {
            this.it = it;
        }

        private T getValue(boolean reset) {
            if (!nextIsCalled) {
                value = it.next();
                nextIsCalled = true;
            }
            if (reset) {
                nextIsCalled = false;
            }
            return value;
        }
    }

    public SortedJoiningIterator(List<Iterator<T>> iterators, Comparator<? super T> elementComparator) {
        int n = iterators.size();
        this.iterators = n == 0 ?
            new PriorityQueue<>() :
            new PriorityQueue<>(n, (h1, h2) -> elementComparator.compare(h1.getValue(false), h2.getValue(false)));
        iterators.stream().filter(Iterator::hasNext).map(Head::new).forEach(this.iterators::add);
    }

    @Override
    public boolean hasNext() {
        return !iterators.isEmpty();
    }

    @Override
    public T next() {
        if (!hasNext()) {
            throw new NoSuchElementException();
        }
        Head<T> head = iterators.poll();
        Iterator<T> it = head.it;
        lastIt = it;
        T value = head.getValue(true);
        if (it.hasNext()) {
            iterators.add(head);
        }
        return value;
    }

    @Override
    public void remove() {
        if (lastIt == null) {
            throw new IllegalStateException();
        }
        lastIt.remove();
    }
}

SortedJoiningIteratorTest.java

package com.exercise;

import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.function.Predicate;
import java.util.stream.Stream;

import static java.util.Collections.emptyIterator;
import static java.util.Comparator.comparingInt;
import static java.util.stream.Collectors.toList;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;

class SortedJoiningIteratorTest {
    @Test
    void noIterators() {
        testEmptyJoiningIterator(List.of());
    }

    @Test
    void oneEmptyIterator() {
        testEmptyJoiningIterator(List.of(emptyIterator()));
    }

    @Test
    void severalEmptyIterators() {
        testEmptyJoiningIterator(List.of(emptyIterator(), emptyIterator(), emptyIterator()));
    }

    @Test
    void oneSingletonIterator() {
        testNotEmptyJoiningIterator(List.of(List.of(1).iterator()), i -> false, List.of(1));
    }

    @Test
    void oneNotSingletonIterator() {
        testNotEmptyJoiningIterator(List.of(List.of(1, 3, 5, 7).iterator()), i -> false, List.of(1, 3, 5, 7));
    }

    @Test
    void severalSingletonIteratorsAsc() {
        testNotEmptyJoiningIterator(List.of(List.of(1).iterator(), List.of(5).iterator(), List.of(6).iterator()), i -> false, List.of(1, 5, 6));
    }

    @Test
    void severalSingletonIteratorsDesc() {
        testNotEmptyJoiningIterator(List.of(List.of(5).iterator(), List.of(3).iterator(), List.of(2).iterator()), i -> false, List.of(2, 3, 5));
    }

    @Test
    void severalNotSingletonIteratorsNoDuplicates() {
        testNotEmptyJoiningIterator(List.of(List.of(1, 3, 5).iterator(), List.of(2, 4, 6).iterator(), List.of(7, 8, 9).iterator()), i -> false, List.of(1, 2, 3, 4, 5, 6, 7, 8, 9));
    }

    @Test
    void severalNotSingletonIteratorsWithDuplicates() {
        testNotEmptyJoiningIterator(List.of(List.of(1, 3, 5).iterator(), List.of(2, 4, 6).iterator(), List.of(2, 3, 4).iterator()), i -> false, List.of(1, 2, 2, 3, 3, 4, 4, 5, 6));
    }

    @Test
    void readOnlyListRemoveAttempt() {
        assertThrows(UnsupportedOperationException.class, () -> testNotEmptyJoiningIterator(List.of(List.of(1, 3, 5).iterator()), i -> i == 3, List.of(1, 2, 2, 4, 4, 5, 6)));
    }

    @Test
    void remove() {
        List<Integer> list = new ArrayList<>(List.of(1,2,3,4));
        testNotEmptyJoiningIterator(List.of(list.iterator()), i -> i == 3, List.of(1, 2, 4));
        assertEquals(List.of(1,2,4), list);
    }

    @Test
    void task1() {
        List<Integer> l1 = List.of(6, 8, 19, 21, 32, 66, 67, 77, 89);
        List<Integer> l2 = List.of(1, 3, 5, 24, 33, 45, 57, 59, 89);
        List<Integer> l3 = List.of(2, 4, 9, 18, 22, 44, 46, 89, 89);
        List<Integer> all = Stream.of(l1, l2, l3).flatMap(Collection::stream).sorted().collect(toList());
        testNotEmptyJoiningIterator(List.of(l1.iterator(), l2.iterator(), l3.iterator()), i -> false, all);
    }

    private void testEmptyJoiningIterator(List<Iterator<Integer>> iterators) {
        Iterator<Integer> it = new SortedJoiningIterator<>(iterators, comparingInt(i -> i));
        assertFalse(it.hasNext());
        assertThrows(NoSuchElementException.class, it::next);
        assertThrows(IllegalStateException.class, it::remove);
    }

    private void testNotEmptyJoiningIterator(List<Iterator<Integer>> iterators, Predicate<Integer> removeFilter, List<Integer> expected) {
        Iterator<Integer> it = new SortedJoiningIterator<>(iterators, comparingInt(i -> i));
        List<Integer> actual = new LinkedList<>();
        while (it.hasNext()) {
            Integer value = it.next();
            if (removeFilter.test(value)) {
                it.remove();
            } else {
                actual.add(value);
            }
        }
        assertEquals(expected, actual);
    }
}