@Test
public void testTokenAware() throws Exception {

    TreeMap<HostToken, HostConnectionPool<Integer>> pools = new TreeMap<HostToken, HostConnectionPool<Integer>>(new Comparator<HostToken>() {
        @Override
        public int compare(HostToken o1, HostToken o2) {
            return o1.getHost().getHostAddress().compareTo(o2.getHost().getHostAddress());
        }
    });
    pools.put(h1, getMockHostConnectionPool(h1));
    pools.put(h2, getMockHostConnectionPool(h2));
    pools.put(h3, getMockHostConnectionPool(h3));
    pools.put(h4, getMockHostConnectionPool(h4));
    TokenAwareSelection<Integer> tokenAwareSelector = new TokenAwareSelection<Integer>();
    tokenAwareSelector.initWithHosts(pools);
    Map<String, Integer> result = new HashMap<String, Integer>();
    runTest(0L, 100000L, result, tokenAwareSelector);
    System.out.println("Token distribution: " + result);
    verifyTokenDistribution(result.values());
}