Unit test coverage for replica map generation

thobbs · thobbs · commit d8b24f8a640e · 2013-09-25T11:11:46.000-05:00
diff --git a/tests/unit/test_metadata.py b/tests/unit/test_metadata.py
@@ -1,9 +1,12 @@
 import unittest
+
 import cassandra
 from cassandra.metadata import (TableMetadata, Murmur3Token, MD5Token,
                                 BytesToken, ReplicationStrategy,
                                 NetworkTopologyStrategy, SimpleStrategy,
                                 LocalStrategy)
+from cassandra.policies import SimpleConvictionPolicy
+from cassandra.pool import Host
 
 
 class TestStrategies(unittest.TestCase):
@@ -50,11 +53,80 @@ def test_replication_strategy(self):
         self.assertRaises(NotImplementedError, rs.make_token_replica_map, None, None)
         self.assertRaises(NotImplementedError, rs.export_for_schema)
 
-    def test_nts(self):
-        # TODO: Cover NetworkTopologyStrategy.make_token_replica_map()
+    def test_nts_make_token_replica_map(self):
+        token_to_host_owner = {}
+
+        dc1_1 = Host('dc1.1', SimpleConvictionPolicy)
+        dc1_2 = Host('dc1.2', SimpleConvictionPolicy)
+        dc1_3 = Host('dc1.3', SimpleConvictionPolicy)
+        for host in (dc1_1, dc1_2, dc1_3):
+            host.set_location_info('dc1', 'rack1')
+        token_to_host_owner[MD5Token(0)] = dc1_1
+        token_to_host_owner[MD5Token(100)] = dc1_2
+        token_to_host_owner[MD5Token(200)] = dc1_3
+
+        dc2_1 = Host('dc2.1', SimpleConvictionPolicy)
+        dc2_2 = Host('dc2.2', SimpleConvictionPolicy)
+        dc2_1.set_location_info('dc2', 'rack1')
+        dc2_2.set_location_info('dc2', 'rack1')
+        token_to_host_owner[MD5Token(1)] = dc2_1
+        token_to_host_owner[MD5Token(101)] = dc2_2
+
+        dc3_1 = Host('dc3.1', SimpleConvictionPolicy)
+        dc3_1.set_location_info('dc3', 'rack3')
+        token_to_host_owner[MD5Token(2)] = dc3_1
+
+        ring = [MD5Token(0),
+                MD5Token(1),
+                MD5Token(2),
+                MD5Token(100),
+                MD5Token(101),
+                MD5Token(200)]
+
+        nts = NetworkTopologyStrategy({'dc1': 2, 'dc2': 2, 'dc3': 1})
+        replica_map = nts.make_token_replica_map(token_to_host_owner, ring)
+
+        self.assertItemsEqual(replica_map[MD5Token(0)], (dc1_1, dc1_2, dc2_1, dc2_2, dc3_1))
+
+    def test_nts_make_token_replica_map_empty_dc(self):
+        host = Host('1', SimpleConvictionPolicy)
+        host.set_location_info('dc1', 'rack1')
+        token_to_host_owner = {MD5Token(0): host}
+        ring = [MD5Token(0)]
+        nts = NetworkTopologyStrategy({'dc1': 1, 'dc2': 0})
+
+        replica_map = nts.make_token_replica_map(token_to_host_owner, ring)
+        self.assertEqual(replica_map[MD5Token(0)], set([host]))
+
+    def test_nts_export_for_schema(self):
         # TODO: Cover NetworkTopologyStrategy.export_for_schema()
         pass
 
+    def test_simple_strategy_make_token_replica_map(self):
+        host1 = Host('1', SimpleConvictionPolicy)
+        host2 = Host('2', SimpleConvictionPolicy)
+        host3 = Host('3', SimpleConvictionPolicy)
+        token_to_host_owner = {
+            MD5Token(0): host1,
+            MD5Token(100): host2,
+            MD5Token(200): host3
+        }
+        ring = [MD5Token(0), MD5Token(100), MD5Token(200)]
+
+        rf1_replicas = SimpleStrategy(1).make_token_replica_map(token_to_host_owner, ring)
+        self.assertItemsEqual(rf1_replicas[MD5Token(0)], [host1])
+        self.assertItemsEqual(rf1_replicas[MD5Token(100)], [host2])
+        self.assertItemsEqual(rf1_replicas[MD5Token(200)], [host3])
+
+        rf2_replicas = SimpleStrategy(2).make_token_replica_map(token_to_host_owner, ring)
+        self.assertItemsEqual(rf2_replicas[MD5Token(0)], [host1, host2])
+        self.assertItemsEqual(rf2_replicas[MD5Token(100)], [host2, host3])
+        self.assertItemsEqual(rf2_replicas[MD5Token(200)], [host3, host1])
+
+        rf3_replicas = SimpleStrategy(3).make_token_replica_map(token_to_host_owner, ring)
+        self.assertItemsEqual(rf3_replicas[MD5Token(0)], [host1, host2, host3])
+        self.assertItemsEqual(rf3_replicas[MD5Token(100)], [host2, host3, host1])
+        self.assertItemsEqual(rf3_replicas[MD5Token(200)], [host3, host1, host2])
 
 class TestTokens(unittest.TestCase):
 
@@ -141,7 +213,7 @@ def test_token_values(self):
         md5_token = MD5Token(cassandra.metadata.MIN_LONG - 1)
         self.assertEqual(md5_token.hash_fn('123'), 42767516990368493138776584305024125808L)
         self.assertEqual(md5_token.hash_fn(str(cassandra.metadata.MAX_LONG)), 28528976619278518853815276204542453639L)
-        self.assertEqual(str(md5_token), '<MD5Token: -9223372036854775809>')
+        self.assertEqual(str(md5_token), '<MD5Token: -9223372036854775809L>')
 
         bytes_token = BytesToken(str(cassandra.metadata.MIN_LONG - 1))
         self.assertEqual(bytes_token.hash_fn('123'), '123')
