### Question:
HostAffinityPartitionStrategy implements PartitioningStrategy { @Override public void distributeTestTestTimes(PartitionRequest request, Collection<TestTime> testTimes, Set<TestContainer> testContainers) { if(isNewHostDetected(request.getHostList(), testContainers)) { LOG.info("New host detected. Performing a full greedy partition"); PartitionUtil.greedyPartitionFailuresOnSameHost(testTimes, testContainers); } else { List<TestTime> testsWithNoHostAffinity = assignTestsWithHostAffinities(testTimes, testContainers); PartitionUtil.greedyPartitionFailuresOnSameHost(testsWithNoHostAffinity, testContainers); Rebalancer rebalancer = Rebalancer.newTimeRebalancer(rebalanceThreshold); if (rebalancer.isRebalanceNeeded(testContainers)) { rebalancer.balanceTestContainers(testContainers); } } }  @Override void distributeTestTestTimes(PartitionRequest request, Collection<TestTime> testTimes, Set<TestContainer> testContainers);  }### Answer:
@Test public void partitionWithHostNames() { String[] hostNames = {"host1", "host2"}; Collection<TestTime> testData = generateTests(10, hostNames); Set<String> hostSet = Stream.of(hostNames).collect(Collectors.toSet()); Set<TestContainer> containers = hostSet.stream().map(h -> new TestContainer(h, "whatever")).collect(Collectors.toSet()); partitioner.distributeTestTestTimes(new PartitionRequest("host1", "whatever", hostSet, "tag"), testData, containers); containers.forEach(c -> assertEquals(5, c.getTestTimes().size())); }
@Test public void partitionWithoutHostNames() { Collection<TestTime> testData = generateTests(10); Set<String> hostSet = Sets.newHashSet("host1", "host2"); Set<TestContainer> containers = hostSet.stream().map(h -> new TestContainer(h, "whatever")).collect(Collectors.toSet()); partitioner.distributeTestTestTimes(new PartitionRequest("host1", "whatever", hostSet, "tag"), testData, containers); containers.forEach(c -> assertEquals(5, c.getTestTimes().size())); }