Added test for WhiteListRoundRobinPolicyTests

beltran · beltran · commit 4ef23f3d7b19 · 2017-10-03T17:03:49.000-04:00
diff --git a/tests/integration/standard/test_policies.py b/tests/integration/standard/test_policies.py
@@ -19,19 +19,19 @@
 except ImportError:
     import unittest  # noqa
 
-from cassandra.cluster import Cluster
-from cassandra.policies import HostFilterPolicy, RoundRobinPolicy, \
-    SimpleConvictionPolicy
+from cassandra.cluster import Cluster, ExecutionProfile
+from cassandra.policies import HostFilterPolicy, RoundRobinPolicy,  SimpleConvictionPolicy, WhiteListRoundRobinPolicy
 from cassandra.pool import Host
 
-from tests.integration import PROTOCOL_VERSION
-
+from tests.integration import PROTOCOL_VERSION, local
 
 from concurrent.futures import wait as wait_futures
 
+
 def setup_module():
     use_singledc()
 
+
 class HostFilterPolicyTests(unittest.TestCase):
 
     def test_predicate_changes(self):
@@ -74,3 +74,20 @@ def test_predicate_changes(self):
             response = session.execute("SELECT * from system.local")
             queried_hosts.update(response.response_future.attempted_hosts)
         self.assertEqual(queried_hosts, all_hosts)
+
+
+class WhiteListRoundRobinPolicyTests(unittest.TestCase):
+
+    @local
+    def test_only_connects_to_subset(self):
+        only_connect_hosts = {"127.0.0.1", "127.0.0.2"}
+        white_list = ExecutionProfile(load_balancing_policy=WhiteListRoundRobinPolicy(only_connect_hosts))
+        cluster = Cluster(execution_profiles={"white_list": white_list})
+        #cluster = Cluster(load_balancing_policy=WhiteListRoundRobinPolicy(only_connect_hosts))
+        session = cluster.connect(wait_for_all_pools=True)
+        queried_hosts = set()
+        for _ in range(10):
+            response = session.execute('SELECT * from system.local', execution_profile="white_list")
+            queried_hosts.update(response.response_future.attempted_hosts)
+        queried_hosts = set(host.address for host in queried_hosts)
+        self.assertEqual(queried_hosts, only_connect_hosts)
diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py
@@ -1232,22 +1232,6 @@ def test_hosts_with_hostname(self):
 
         self.assertEqual(policy.distance(host), HostDistance.LOCAL)
 
-    def test_deprecated(self):
-        import warnings
-
-        warnings.resetwarnings()  # in case we've instantiated one before
-
-        # set up warning filters to allow all, set up restore when this test is done
-        filters_backup, warnings.filters = warnings.filters, []
-        self.addCleanup(setattr, warnings, 'filters', filters_backup)
-
-        with warnings.catch_warnings(record=True) as caught_warnings:
-            WhiteListRoundRobinPolicy([])
-            self.assertEqual(len(caught_warnings), 1)
-            warning_message = caught_warnings[-1]
-            self.assertEqual(warning_message.category, DeprecationWarning)
-            self.assertIn('4.0', warning_message.message.args[0])
-
 
 class AddressTranslatorTest(unittest.TestCase):
 
