From 08a2370157d33f0aff51aabe156dc0f6a74ad7f5 Mon Sep 17 00:00:00 2001 From: Tyler Hobbs Date: Wed, 3 Jul 2013 12:47:55 -0500 Subject: [PATCH] Update DCAwareLoadBalancingPolicy unit tests --- tests/unit/test_policies.py | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py index 7a6c29fe..a87f881f 100644 --- a/tests/unit/test_policies.py +++ b/tests/unit/test_policies.py @@ -78,11 +78,32 @@ class TestDCAwareRoundRobinPolicy(unittest.TestCase): for h in hosts[2:]: h.set_location_info("dc2", "rack1") - policy = DCAwareRoundRobinPolicy("dc1") + local_hosts = set(h for h in hosts if h.datacenter == "dc1") + remote_hosts = set(h for h in hosts if h.datacenter != "dc1") + + # allow all of the remote hosts to be used + policy = DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=2) policy.populate(None, hosts) qplan = list(policy.make_query_plan()) - self.assertEqual(set(qplan[:2]), set(h for h in hosts if h.datacenter == "dc1")) - self.assertEqual(set(qplan[2:]), set(h for h in hosts if h.datacenter != "dc1")) + self.assertEqual(set(qplan[:2]), local_hosts) + self.assertEqual(set(qplan[2:]), remote_hosts) + + # allow only one of the remote hosts to be used + policy = DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=1) + policy.populate(None, hosts) + qplan = list(policy.make_query_plan()) + self.assertEqual(set(qplan[:2]), local_hosts) + + used_remotes = set(qplan[2:]) + self.assertEqual(1, len(used_remotes)) + self.assertIn(qplan[2], remote_hosts) + + # allow no remote hosts to be used + policy = DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=0) + policy.populate(None, hosts) + qplan = list(policy.make_query_plan()) + self.assertEqual(2, len(qplan)) + self.assertEqual(local_hosts, set(qplan)) def test_get_distance(self): policy = DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=0)