Code review comment: These 2 can be updated to hasValue as well.

Code before review:
void testConcurrentSplitAssignmentForMultipleHosts() throws InterruptedException {
final int NUM_THREADS = 10;
final int NUM_SPLITS = 500;
final int SUM_OF_IDS = (NUM_SPLITS - 1) * (NUM_SPLITS) / 2;
final String[] hosts = {"host1", "host1", "host1", "host2", "host2", "host3"};
Set<LocatableInputSplit> splits = new HashSet<>();
for (int i = 0; i < NUM_SPLITS; i++) {
splits.add(new LocatableInputSplit(i, hosts[i % hosts.length]));
}
final LocatableInputSplitAssigner ia = new LocatableInputSplitAssigner(splits);
final AtomicInteger splitsRetrieved = new AtomicInteger(0);
final AtomicInteger sumOfIds = new AtomicInteger(0);
Runnable retriever =
() -> {
final String threadHost = hosts[(int) (Math.random() * hosts.length)];
LocatableInputSplit split;
while ((split = ia.getNextInputSplit(threadHost, 0)) != null) {
splitsRetrieved.incrementAndGet();
sumOfIds.addAndGet(split.getSplitNumber());
}
};
Thread[] threads = new Thread[NUM_THREADS];
for (int i = 0; i < NUM_THREADS; i++) {
threads[i] = new Thread(retriever);
threads[i].setDaemon(true);
}
for (int i = 0; i < NUM_THREADS; i++) {
threads[i].start();
}
for (int i = 0; i < NUM_THREADS; i++) {
threads[i].join(5000);
}
for (int i = 0; i < NUM_THREADS; i++) {
assertThat(threads[i].isAlive()).isFalse();
}
assertThat(splitsRetrieved.get()).isEqualTo(NUM_SPLITS);
assertThat(sumOfIds.get()).isEqualTo(SUM_OF_IDS);
assertThat(ia.getNextInputSplit("testhost", 0)).isNull();
assertThat(ia.getNumberOfRemoteAssignments())
.isGreaterThanOrEqualTo(NUM_SPLITS / hosts.length);
}

Line under review: assertThat(sumOfIds.get()).isEqualTo(SUM_OF_IDS);

Code after addressing the review:
void testConcurrentSplitAssignmentForMultipleHosts() throws InterruptedException {
final int NUM_THREADS = 10;
final int NUM_SPLITS = 500;
final int SUM_OF_IDS = (NUM_SPLITS - 1) * (NUM_SPLITS) / 2;
final String[] hosts = {"host1", "host1", "host1", "host2", "host2", "host3"};
Set<LocatableInputSplit> splits = new HashSet<>();
for (int i = 0; i < NUM_SPLITS; i++) {
splits.add(new LocatableInputSplit(i, hosts[i % hosts.length]));
}
final LocatableInputSplitAssigner ia = new LocatableInputSplitAssigner(splits);
final AtomicInteger splitsRetrieved = new AtomicInteger(0);
final AtomicInteger sumOfIds = new AtomicInteger(0);
Runnable retriever =
() -> {
final String threadHost = hosts[(int) (Math.random() * hosts.length)];
LocatableInputSplit split;
while ((split = ia.getNextInputSplit(threadHost, 0)) != null) {
splitsRetrieved.incrementAndGet();
sumOfIds.addAndGet(split.getSplitNumber());
}
};
Thread[] threads = new Thread[NUM_THREADS];
for (int i = 0; i < NUM_THREADS; i++) {
threads[i] = new Thread(retriever);
threads[i].setDaemon(true);
}
for (int i = 0; i < NUM_THREADS; i++) {
threads[i].start();
}
for (int i = 0; i < NUM_THREADS; i++) {
threads[i].join(5000);
}
for (int i = 0; i < NUM_THREADS; i++) {
assertThat(threads[i].isAlive()).isFalse();
}
assertThat(splitsRetrieved).hasValue(NUM_SPLITS);
assertThat(sumOfIds).hasValue(SUM_OF_IDS);
assertThat(ia.getNextInputSplit("testhost", 0)).isNull();
assertThat(ia.getNumberOfRemoteAssignments())
.isGreaterThanOrEqualTo(NUM_SPLITS / hosts.length);
}