package com.netflix.astyanax.connectionpool.impl; import java.math.BigInteger; import java.util.Arrays; import java.util.Collection; import java.util.List; import java.util.Map; import junit.framework.Assert; import org.junit.Test; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.netflix.astyanax.connectionpool.Host; import com.netflix.astyanax.connectionpool.HostConnectionPool; import com.netflix.astyanax.connectionpool.LatencyScoreStrategy; import com.netflix.astyanax.test.TestClient; import com.netflix.astyanax.test.TestHostConnectionPool; public class HostConnectionPoolPartitionTest { @Test public void testPartition() { LatencyScoreStrategy strategy = new SmaLatencyScoreStrategyImpl(10000, 60000, 100, 4.0); HostConnectionPoolPartition partition = new HostConnectionPoolPartition( new BigInteger("1"), strategy); List<TestHostConnectionPool> pools = Arrays.asList(makePool(1), makePool(2), makePool(3), makePool(4)); List<TestHostConnectionPool> pool1 = Arrays.asList(pools.get(0), pools.get(1), pools.get(2)); List<TestHostConnectionPool> pool2 = Arrays.asList(pools.get(0), pools.get(3)); Assert.assertEquals(0, partition.getPools().size()); partition.setPools(pool1); Assert.assertEquals(3, partition.getPools().size()); partition.setPools(pool2); Assert.assertEquals(2, partition.getPools().size()); } // Is there a reason that this test is in HostConnectionPoolPartitionTest? // Perhaps this should be moved to the existing TokenAwareConnectionPoolTest or // a new TokenParitionedTopologyTest? @Test public void testTopology() { LatencyScoreStrategy strategy = new SmaLatencyScoreStrategyImpl(10000, 60000, 100, 4.0); int nHosts = 6; int nReplicationFactor = 3; TokenPartitionedTopology<TestClient> topology = new TokenPartitionedTopology<TestClient>( strategy); // Make the set of pools List<HostConnectionPool<TestClient>> pools = Lists.newArrayList(); for (int i = 0; i < nHosts; i++) { pools.add(makePool(i)); } // Make the flat ring (before ring_describe is called) Map<BigInteger, Collection<HostConnectionPool<TestClient>>> flatRing = Maps .newHashMap(); flatRing.put(new BigInteger("0"), pools); boolean didChange = topology.setPools(flatRing); Assert.assertTrue(didChange); Assert.assertEquals(1, topology.getPartitionCount()); System.out.println(topology); // Make a ring with tokens and RF Map<BigInteger, Collection<HostConnectionPool<TestClient>>> tokenRing = Maps .newHashMap(); for (int i = 0; i < nHosts; i++) { List<HostConnectionPool<TestClient>> partition = Lists .newArrayList(); for (int j = 0; j < nReplicationFactor; j++) { partition.add(pools.get((i + j) % nHosts)); } tokenRing .put(new BigInteger(Integer.toString(i * 1000)), partition); } didChange = topology.setPools(tokenRing); Assert.assertTrue(didChange); Assert.assertEquals(nHosts, topology.getPartitionCount()); System.out.println(topology); HostConnectionPoolPartition<TestClient> partition = topology .getPartition(null); Assert.assertEquals(nHosts, partition.getPools().size()); // Partition Token Map: // HCP 0 - (5000, 0] // HCP 1000 - (0, 1000] // HCP 2000 - (1000, 2000] // HCP 3000 - (2000, 3000] // HCP 4000 - (3000, 4000] // HCP 5000 - (4000, 5000] // Test ordinals for (int i = 0; i < nHosts; i++) { partition = topology.getPartition(new BigInteger(Integer .toString(i * 1000))); Assert.assertEquals(new BigInteger(Integer.toString(i * 1000)), partition.id()); } // Test mid-range tokens for (int i = nHosts; i > 0; i--) { partition = topology.getPartition(new BigInteger(Integer .toString(i * 1000 - 500))); if (i == nHosts) { // 5500 is contained in (5000,0] which belongs to HCP 0 Assert.assertEquals(BigInteger.ZERO, partition.id()); } else { Assert.assertEquals(new BigInteger(Integer.toString(i * 1000)), partition.id()); } } Map<BigInteger, Collection<HostConnectionPool<TestClient>>> emptyRing = Maps .newHashMap(); topology.setPools(emptyRing); System.out.println(topology); Assert.assertEquals(0, topology.getPartitionCount()); Assert.assertEquals(0, topology.getAllPools().getPools().size()); } public TestHostConnectionPool makePool(int index) { return new TestHostConnectionPool(new Host("127.0.0." + index, 0)); } }