# Copyright DataStax, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest from itertools import islice, cycle from unittest.mock import Mock, patch, call from random import randint from _thread import LockType import sys import struct from threading import Thread from cassandra import ConsistencyLevel from cassandra.cluster import Cluster from cassandra.connection import DefaultEndPoint from cassandra.metadata import Metadata from cassandra.policies import (RoundRobinPolicy, WhiteListRoundRobinPolicy, DCAwareRoundRobinPolicy, TokenAwarePolicy, SimpleConvictionPolicy, HostDistance, ExponentialReconnectionPolicy, RetryPolicy, WriteType, DowngradingConsistencyRetryPolicy, ConstantReconnectionPolicy, LoadBalancingPolicy, ConvictionPolicy, ReconnectionPolicy, FallthroughRetryPolicy, IdentityTranslator, EC2MultiRegionTranslator, HostFilterPolicy) from cassandra.pool import Host from cassandra.query import Statement class LoadBalancingPolicyTest(unittest.TestCase): def test_non_implemented(self): """ Code coverage for interface-style base class """ policy = LoadBalancingPolicy() host = Host(DefaultEndPoint("ip1"), SimpleConvictionPolicy) host.set_location_info("dc1", "rack1") self.assertRaises(NotImplementedError, policy.distance, host) self.assertRaises(NotImplementedError, policy.populate, None, host) self.assertRaises(NotImplementedError, policy.make_query_plan) self.assertRaises(NotImplementedError, policy.on_up, host) self.assertRaises(NotImplementedError, policy.on_down, host) self.assertRaises(NotImplementedError, policy.on_add, host) self.assertRaises(NotImplementedError, policy.on_remove, host) def test_instance_check(self): self.assertRaises(TypeError, Cluster, load_balancing_policy=RoundRobinPolicy) class RoundRobinPolicyTest(unittest.TestCase): def test_basic(self): hosts = [0, 1, 2, 3] policy = RoundRobinPolicy() policy.populate(None, hosts) qplan = list(policy.make_query_plan()) self.assertEqual(sorted(qplan), hosts) def test_multiple_query_plans(self): hosts = [0, 1, 2, 3] policy = RoundRobinPolicy() policy.populate(None, hosts) for i in range(20): qplan = list(policy.make_query_plan()) self.assertEqual(sorted(qplan), hosts) def test_single_host(self): policy = RoundRobinPolicy() policy.populate(None, [0]) qplan = list(policy.make_query_plan()) self.assertEqual(qplan, [0]) def test_status_updates(self): hosts = [0, 1, 2, 3] policy = RoundRobinPolicy() policy.populate(None, hosts) policy.on_down(0) policy.on_remove(1) policy.on_up(4) policy.on_add(5) qplan = list(policy.make_query_plan()) self.assertEqual(sorted(qplan), [2, 3, 4, 5]) def test_thread_safety(self): hosts = range(100) policy = RoundRobinPolicy() policy.populate(None, hosts) def check_query_plan(): for i in range(100): qplan = list(policy.make_query_plan()) self.assertEqual(sorted(qplan), list(hosts)) threads = [Thread(target=check_query_plan) for i in range(4)] for t in threads: t.start() for t in threads: t.join() def test_thread_safety_during_modification(self): hosts = range(100) policy = RoundRobinPolicy() policy.populate(None, hosts) errors = [] def check_query_plan(): try: for i in range(100): list(policy.make_query_plan()) except Exception as exc: errors.append(exc) def host_up(): for i in range(1000): policy.on_up(randint(0, 99)) def host_down(): for i in range(1000): policy.on_down(randint(0, 99)) threads = [] for i in range(5): threads.append(Thread(target=check_query_plan)) threads.append(Thread(target=host_up)) threads.append(Thread(target=host_down)) # make the GIL switch after every instruction, maximizing # the chance of race conditions check = '__pypy__' in sys.builtin_module_names if check: original_interval = sys.getcheckinterval() else: original_interval = sys.getswitchinterval() try: if check: sys.setcheckinterval(0) else: sys.setswitchinterval(0.0001) for t in threads: t.start() for t in threads: t.join() finally: if check: sys.setcheckinterval(original_interval) else: sys.setswitchinterval(original_interval) if errors: self.fail("Saw errors: %s" % (errors,)) def test_no_live_nodes(self): """ Ensure query plan for a downed cluster will execute without errors """ hosts = [0, 1, 2, 3] policy = RoundRobinPolicy() policy.populate(None, hosts) for i in range(4): policy.on_down(i) qplan = list(policy.make_query_plan()) self.assertEqual(qplan, []) class DCAwareRoundRobinPolicyTest(unittest.TestCase): def test_no_remote(self): hosts = [] for i in range(4): h = Host(DefaultEndPoint(i), SimpleConvictionPolicy) h.set_location_info("dc1", "rack1") hosts.append(h) policy = DCAwareRoundRobinPolicy("dc1") policy.populate(None, hosts) qplan = list(policy.make_query_plan()) self.assertEqual(sorted(qplan), sorted(hosts)) def test_with_remotes(self): hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy) for i in range(4)] for h in hosts[:2]: h.set_location_info("dc1", "rack1") for h in hosts[2:]: h.set_location_info("dc2", "rack1") 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(Mock(), hosts) qplan = list(policy.make_query_plan()) 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(Mock(), 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(Mock(), 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) host = Host(DefaultEndPoint("ip1"), SimpleConvictionPolicy) host.set_location_info("dc1", "rack1") policy.populate(Mock(), [host]) self.assertEqual(policy.distance(host), HostDistance.LOCAL) # used_hosts_per_remote_dc is set to 0, so ignore it remote_host = Host(DefaultEndPoint("ip2"), SimpleConvictionPolicy) remote_host.set_location_info("dc2", "rack1") self.assertEqual(policy.distance(remote_host), HostDistance.IGNORED) # dc2 isn't registered in the policy's live_hosts dict policy.used_hosts_per_remote_dc = 1 self.assertEqual(policy.distance(remote_host), HostDistance.IGNORED) # make sure the policy has both dcs registered policy.populate(Mock(), [host, remote_host]) self.assertEqual(policy.distance(remote_host), HostDistance.REMOTE) # since used_hosts_per_remote_dc is set to 1, only the first # remote host in dc2 will be REMOTE, the rest are IGNORED second_remote_host = Host(DefaultEndPoint("ip3"), SimpleConvictionPolicy) second_remote_host.set_location_info("dc2", "rack1") policy.populate(Mock(), [host, remote_host, second_remote_host]) distances = set([policy.distance(remote_host), policy.distance(second_remote_host)]) self.assertEqual(distances, set([HostDistance.REMOTE, HostDistance.IGNORED])) def test_status_updates(self): hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy) for i in range(4)] for h in hosts[:2]: h.set_location_info("dc1", "rack1") for h in hosts[2:]: h.set_location_info("dc2", "rack1") policy = DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=1) policy.populate(Mock(), hosts) policy.on_down(hosts[0]) policy.on_remove(hosts[2]) new_local_host = Host(DefaultEndPoint(4), SimpleConvictionPolicy) new_local_host.set_location_info("dc1", "rack1") policy.on_up(new_local_host) new_remote_host = Host(DefaultEndPoint(5), SimpleConvictionPolicy) new_remote_host.set_location_info("dc9000", "rack1") policy.on_add(new_remote_host) # we now have two local hosts and two remote hosts in separate dcs qplan = list(policy.make_query_plan()) self.assertEqual(set(qplan[:2]), set([hosts[1], new_local_host])) self.assertEqual(set(qplan[2:]), set([hosts[3], new_remote_host])) # since we have hosts in dc9000, the distance shouldn't be IGNORED self.assertEqual(policy.distance(new_remote_host), HostDistance.REMOTE) policy.on_down(new_local_host) policy.on_down(hosts[1]) qplan = list(policy.make_query_plan()) self.assertEqual(set(qplan), set([hosts[3], new_remote_host])) policy.on_down(new_remote_host) policy.on_down(hosts[3]) qplan = list(policy.make_query_plan()) self.assertEqual(qplan, []) def test_modification_during_generation(self): hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy) for i in range(4)] for h in hosts[:2]: h.set_location_info("dc1", "rack1") for h in hosts[2:]: h.set_location_info("dc2", "rack1") policy = DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=3) policy.populate(Mock(), hosts) # The general concept here is to change thee internal state of the # policy during plan generation. In this case we use a grey-box # approach that changes specific things during known phases of the # generator. new_host = Host(DefaultEndPoint(4), SimpleConvictionPolicy) new_host.set_location_info("dc1", "rack1") # new local before iteration plan = policy.make_query_plan() policy.on_up(new_host) # local list is not bound yet, so we get to see that one self.assertEqual(len(list(plan)), 3 + 2) # remove local before iteration plan = policy.make_query_plan() policy.on_down(new_host) # local list is not bound yet, so we don't see it self.assertEqual(len(list(plan)), 2 + 2) # new local after starting iteration plan = policy.make_query_plan() next(plan) policy.on_up(new_host) # local list was is bound, and one consumed, so we only see the other original self.assertEqual(len(list(plan)), 1 + 2) # remove local after traversing available plan = policy.make_query_plan() for _ in range(3): next(plan) policy.on_down(new_host) # we should be past the local list self.assertEqual(len(list(plan)), 0 + 2) # REMOTES CHANGE new_host.set_location_info("dc2", "rack1") # new remote after traversing local, but not starting remote plan = policy.make_query_plan() for _ in range(2): next(plan) policy.on_up(new_host) # list is updated before we get to it self.assertEqual(len(list(plan)), 0 + 3) # remove remote after traversing local, but not starting remote plan = policy.make_query_plan() for _ in range(2): next(plan) policy.on_down(new_host) # list is updated before we get to it self.assertEqual(len(list(plan)), 0 + 2) # new remote after traversing local, and starting remote plan = policy.make_query_plan() for _ in range(3): next(plan) policy.on_up(new_host) # slice is already made, and we've consumed one self.assertEqual(len(list(plan)), 0 + 1) # remove remote after traversing local, and starting remote plan = policy.make_query_plan() for _ in range(3): next(plan) policy.on_down(new_host) # slice is created with all present, and we've consumed one self.assertEqual(len(list(plan)), 0 + 2) # local DC disappears after finishing it, but not starting remote plan = policy.make_query_plan() for _ in range(2): next(plan) policy.on_down(hosts[0]) policy.on_down(hosts[1]) # dict traversal starts as normal self.assertEqual(len(list(plan)), 0 + 2) policy.on_up(hosts[0]) policy.on_up(hosts[1]) # PYTHON-297 addresses the following cases, where DCs come and go # during generation # local DC disappears after finishing it, and starting remote plan = policy.make_query_plan() for _ in range(3): next(plan) policy.on_down(hosts[0]) policy.on_down(hosts[1]) # dict traversal has begun and consumed one self.assertEqual(len(list(plan)), 0 + 1) policy.on_up(hosts[0]) policy.on_up(hosts[1]) # remote DC disappears after finishing local, but not starting remote plan = policy.make_query_plan() for _ in range(2): next(plan) policy.on_down(hosts[2]) policy.on_down(hosts[3]) # nothing left self.assertEqual(len(list(plan)), 0 + 0) policy.on_up(hosts[2]) policy.on_up(hosts[3]) # remote DC disappears while traversing it plan = policy.make_query_plan() for _ in range(3): next(plan) policy.on_down(hosts[2]) policy.on_down(hosts[3]) # we continue with remainder of original list self.assertEqual(len(list(plan)), 0 + 1) policy.on_up(hosts[2]) policy.on_up(hosts[3]) another_host = Host(DefaultEndPoint(5), SimpleConvictionPolicy) another_host.set_location_info("dc3", "rack1") new_host.set_location_info("dc3", "rack1") # new DC while traversing remote plan = policy.make_query_plan() for _ in range(3): next(plan) policy.on_up(new_host) policy.on_up(another_host) # we continue with remainder of original list self.assertEqual(len(list(plan)), 0 + 1) # remote DC disappears after finishing it plan = policy.make_query_plan() for _ in range(3): next(plan) last_host_in_this_dc = next(plan) if last_host_in_this_dc in (new_host, another_host): down_hosts = [new_host, another_host] else: down_hosts = hosts[2:] for h in down_hosts: policy.on_down(h) # the last DC has two self.assertEqual(len(list(plan)), 0 + 2) def test_no_live_nodes(self): """ Ensure query plan for a downed cluster will execute without errors """ hosts = [] for i in range(4): h = Host(DefaultEndPoint(i), SimpleConvictionPolicy) h.set_location_info("dc1", "rack1") hosts.append(h) policy = DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=1) policy.populate(Mock(), hosts) for host in hosts: policy.on_down(host) qplan = list(policy.make_query_plan()) self.assertEqual(qplan, []) def test_no_nodes(self): """ Ensure query plan for an empty cluster will execute without errors """ policy = DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=1) policy.populate(None, []) qplan = list(policy.make_query_plan()) self.assertEqual(qplan, []) def test_default_dc(self): host_local = Host(DefaultEndPoint(1), SimpleConvictionPolicy, 'local') host_remote = Host(DefaultEndPoint(2), SimpleConvictionPolicy, 'remote') host_none = Host(DefaultEndPoint(1), SimpleConvictionPolicy) # contact point is '1' cluster = Mock(endpoints_resolved=[DefaultEndPoint(1)]) # contact DC first policy = DCAwareRoundRobinPolicy() policy.populate(cluster, [host_none]) self.assertFalse(policy.local_dc) policy.on_add(host_local) policy.on_add(host_remote) self.assertNotEqual(policy.local_dc, host_remote.datacenter) self.assertEqual(policy.local_dc, host_local.datacenter) # contact DC second policy = DCAwareRoundRobinPolicy() policy.populate(cluster, [host_none]) self.assertFalse(policy.local_dc) policy.on_add(host_remote) policy.on_add(host_local) self.assertNotEqual(policy.local_dc, host_remote.datacenter) self.assertEqual(policy.local_dc, host_local.datacenter) # no DC policy = DCAwareRoundRobinPolicy() policy.populate(cluster, [host_none]) self.assertFalse(policy.local_dc) policy.on_add(host_none) self.assertFalse(policy.local_dc) # only other DC policy = DCAwareRoundRobinPolicy() policy.populate(cluster, [host_none]) self.assertFalse(policy.local_dc) policy.on_add(host_remote) self.assertFalse(policy.local_dc) class TokenAwarePolicyTest(unittest.TestCase): def test_wrap_round_robin(self): cluster = Mock(spec=Cluster) cluster.metadata = Mock(spec=Metadata) hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy) for i in range(4)] for host in hosts: host.set_up() def get_replicas(keyspace, packed_key): index = struct.unpack('>i', packed_key)[0] return list(islice(cycle(hosts), index, index + 2)) cluster.metadata.get_replicas.side_effect = get_replicas policy = TokenAwarePolicy(RoundRobinPolicy()) policy.populate(cluster, hosts) for i in range(4): query = Statement(routing_key=struct.pack('>i', i), keyspace='keyspace_name') qplan = list(policy.make_query_plan(None, query)) replicas = get_replicas(None, struct.pack('>i', i)) other = set(h for h in hosts if h not in replicas) self.assertEqual(replicas, qplan[:2]) self.assertEqual(other, set(qplan[2:])) # Should use the secondary policy for i in range(4): qplan = list(policy.make_query_plan()) self.assertEqual(set(qplan), set(hosts)) def test_wrap_dc_aware(self): cluster = Mock(spec=Cluster) cluster.metadata = Mock(spec=Metadata) hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy) for i in range(4)] for host in hosts: host.set_up() for h in hosts[:2]: h.set_location_info("dc1", "rack1") for h in hosts[2:]: h.set_location_info("dc2", "rack1") def get_replicas(keyspace, packed_key): index = struct.unpack('>i', packed_key)[0] # return one node from each DC if index % 2 == 0: return [hosts[0], hosts[2]] else: return [hosts[1], hosts[3]] cluster.metadata.get_replicas.side_effect = get_replicas policy = TokenAwarePolicy(DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=1)) policy.populate(cluster, hosts) for i in range(4): query = Statement(routing_key=struct.pack('>i', i), keyspace='keyspace_name') qplan = list(policy.make_query_plan(None, query)) replicas = get_replicas(None, struct.pack('>i', i)) # first should be the only local replica self.assertIn(qplan[0], replicas) self.assertEqual(qplan[0].datacenter, "dc1") # then the local non-replica self.assertNotIn(qplan[1], replicas) self.assertEqual(qplan[1].datacenter, "dc1") # then one of the remotes (used_hosts_per_remote_dc is 1, so we # shouldn't see two remotes) self.assertEqual(qplan[2].datacenter, "dc2") self.assertEqual(3, len(qplan)) class FakeCluster: def __init__(self): self.metadata = Mock(spec=Metadata) def test_get_distance(self): """ Same test as DCAwareRoundRobinPolicyTest.test_get_distance() Except a FakeCluster is needed for the metadata variable and policy.child_policy is needed to change child policy settings """ policy = TokenAwarePolicy(DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=0)) host = Host(DefaultEndPoint("ip1"), SimpleConvictionPolicy) host.set_location_info("dc1", "rack1") policy.populate(self.FakeCluster(), [host]) self.assertEqual(policy.distance(host), HostDistance.LOCAL) # used_hosts_per_remote_dc is set to 0, so ignore it remote_host = Host(DefaultEndPoint("ip2"), SimpleConvictionPolicy) remote_host.set_location_info("dc2", "rack1") self.assertEqual(policy.distance(remote_host), HostDistance.IGNORED) # dc2 isn't registered in the policy's live_hosts dict policy._child_policy.used_hosts_per_remote_dc = 1 self.assertEqual(policy.distance(remote_host), HostDistance.IGNORED) # make sure the policy has both dcs registered policy.populate(self.FakeCluster(), [host, remote_host]) self.assertEqual(policy.distance(remote_host), HostDistance.REMOTE) # since used_hosts_per_remote_dc is set to 1, only the first # remote host in dc2 will be REMOTE, the rest are IGNORED second_remote_host = Host(DefaultEndPoint("ip3"), SimpleConvictionPolicy) second_remote_host.set_location_info("dc2", "rack1") policy.populate(self.FakeCluster(), [host, remote_host, second_remote_host]) distances = set([policy.distance(remote_host), policy.distance(second_remote_host)]) self.assertEqual(distances, set([HostDistance.REMOTE, HostDistance.IGNORED])) def test_status_updates(self): """ Same test as DCAwareRoundRobinPolicyTest.test_status_updates() """ hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy) for i in range(4)] for h in hosts[:2]: h.set_location_info("dc1", "rack1") for h in hosts[2:]: h.set_location_info("dc2", "rack1") policy = TokenAwarePolicy(DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=1)) policy.populate(self.FakeCluster(), hosts) policy.on_down(hosts[0]) policy.on_remove(hosts[2]) new_local_host = Host(DefaultEndPoint(4), SimpleConvictionPolicy) new_local_host.set_location_info("dc1", "rack1") policy.on_up(new_local_host) new_remote_host = Host(DefaultEndPoint(5), SimpleConvictionPolicy) new_remote_host.set_location_info("dc9000", "rack1") policy.on_add(new_remote_host) # we now have two local hosts and two remote hosts in separate dcs qplan = list(policy.make_query_plan()) self.assertEqual(set(qplan[:2]), set([hosts[1], new_local_host])) self.assertEqual(set(qplan[2:]), set([hosts[3], new_remote_host])) # since we have hosts in dc9000, the distance shouldn't be IGNORED self.assertEqual(policy.distance(new_remote_host), HostDistance.REMOTE) policy.on_down(new_local_host) policy.on_down(hosts[1]) qplan = list(policy.make_query_plan()) self.assertEqual(set(qplan), set([hosts[3], new_remote_host])) policy.on_down(new_remote_host) policy.on_down(hosts[3]) qplan = list(policy.make_query_plan()) self.assertEqual(qplan, []) def test_statement_keyspace(self): hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy) for i in range(4)] for host in hosts: host.set_up() cluster = Mock(spec=Cluster) cluster.metadata = Mock(spec=Metadata) replicas = hosts[2:] cluster.metadata.get_replicas.return_value = replicas child_policy = Mock() child_policy.make_query_plan.return_value = hosts child_policy.distance.return_value = HostDistance.LOCAL policy = TokenAwarePolicy(child_policy) policy.populate(cluster, hosts) # no keyspace, child policy is called keyspace = None routing_key = 'routing_key' query = Statement(routing_key=routing_key) qplan = list(policy.make_query_plan(keyspace, query)) self.assertEqual(hosts, qplan) self.assertEqual(cluster.metadata.get_replicas.call_count, 0) child_policy.make_query_plan.assert_called_once_with(keyspace, query) # working keyspace, no statement cluster.metadata.get_replicas.reset_mock() keyspace = 'working_keyspace' routing_key = 'routing_key' query = Statement(routing_key=routing_key) qplan = list(policy.make_query_plan(keyspace, query)) self.assertEqual(replicas + hosts[:2], qplan) cluster.metadata.get_replicas.assert_called_with(keyspace, routing_key) # statement keyspace, no working cluster.metadata.get_replicas.reset_mock() working_keyspace = None statement_keyspace = 'statement_keyspace' routing_key = 'routing_key' query = Statement(routing_key=routing_key, keyspace=statement_keyspace) qplan = list(policy.make_query_plan(working_keyspace, query)) self.assertEqual(replicas + hosts[:2], qplan) cluster.metadata.get_replicas.assert_called_with(statement_keyspace, routing_key) # both keyspaces set, statement keyspace used for routing cluster.metadata.get_replicas.reset_mock() working_keyspace = 'working_keyspace' statement_keyspace = 'statement_keyspace' routing_key = 'routing_key' query = Statement(routing_key=routing_key, keyspace=statement_keyspace) qplan = list(policy.make_query_plan(working_keyspace, query)) self.assertEqual(replicas + hosts[:2], qplan) cluster.metadata.get_replicas.assert_called_with(statement_keyspace, routing_key) def test_shuffles_if_given_keyspace_and_routing_key(self): """ Test to validate the hosts are shuffled when `shuffle_replicas` is truthy @since 3.8 @jira_ticket PYTHON-676 @expected_result shuffle should be called, because the keyspace and the routing key are set @test_category policy """ self._assert_shuffle(keyspace='keyspace', routing_key='routing_key') def test_no_shuffle_if_given_no_keyspace(self): """ Test to validate the hosts are not shuffled when no keyspace is provided @since 3.8 @jira_ticket PYTHON-676 @expected_result shuffle should be called, because keyspace is None @test_category policy """ self._assert_shuffle(keyspace=None, routing_key='routing_key') def test_no_shuffle_if_given_no_routing_key(self): """ Test to validate the hosts are not shuffled when no routing_key is provided @since 3.8 @jira_ticket PYTHON-676 @expected_result shuffle should be called, because routing_key is None @test_category policy """ self._assert_shuffle(keyspace='keyspace', routing_key=None) @patch('cassandra.policies.shuffle') def _assert_shuffle(self, patched_shuffle, keyspace, routing_key): hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy) for i in range(4)] for host in hosts: host.set_up() cluster = Mock(spec=Cluster) cluster.metadata = Mock(spec=Metadata) replicas = hosts[2:] cluster.metadata.get_replicas.return_value = replicas child_policy = Mock() child_policy.make_query_plan.return_value = hosts child_policy.distance.return_value = HostDistance.LOCAL policy = TokenAwarePolicy(child_policy, shuffle_replicas=True) policy.populate(cluster, hosts) cluster.metadata.get_replicas.reset_mock() child_policy.make_query_plan.reset_mock() query = Statement(routing_key=routing_key) qplan = list(policy.make_query_plan(keyspace, query)) if keyspace is None or routing_key is None: self.assertEqual(hosts, qplan) self.assertEqual(cluster.metadata.get_replicas.call_count, 0) child_policy.make_query_plan.assert_called_once_with(keyspace, query) self.assertEqual(patched_shuffle.call_count, 0) else: self.assertEqual(set(replicas), set(qplan[:2])) self.assertEqual(hosts[:2], qplan[2:]) child_policy.make_query_plan.assert_called_once_with(keyspace, query) self.assertEqual(patched_shuffle.call_count, 1) class ConvictionPolicyTest(unittest.TestCase): def test_not_implemented(self): """ Code coverage for interface-style base class """ conviction_policy = ConvictionPolicy(1) self.assertRaises(NotImplementedError, conviction_policy.add_failure, 1) self.assertRaises(NotImplementedError, conviction_policy.reset) class SimpleConvictionPolicyTest(unittest.TestCase): def test_basic_responses(self): """ Code coverage for SimpleConvictionPolicy """ conviction_policy = SimpleConvictionPolicy(1) self.assertEqual(conviction_policy.add_failure(1), True) self.assertEqual(conviction_policy.reset(), None) class ReconnectionPolicyTest(unittest.TestCase): def test_basic_responses(self): """ Code coverage for interface-style base class """ policy = ReconnectionPolicy() self.assertRaises(NotImplementedError, policy.new_schedule) class ConstantReconnectionPolicyTest(unittest.TestCase): def test_bad_vals(self): """ Test initialization values """ self.assertRaises(ValueError, ConstantReconnectionPolicy, -1, 0) def test_schedule(self): """ Test ConstantReconnectionPolicy schedule """ delay = 2 max_attempts = 100 policy = ConstantReconnectionPolicy(delay=delay, max_attempts=max_attempts) schedule = list(policy.new_schedule()) self.assertEqual(len(schedule), max_attempts) for i, delay in enumerate(schedule): self.assertEqual(delay, delay) def test_schedule_negative_max_attempts(self): """ Test how negative max_attempts are handled """ delay = 2 max_attempts = -100 try: ConstantReconnectionPolicy(delay=delay, max_attempts=max_attempts) self.fail('max_attempts should throw ValueError when negative') except ValueError: pass def test_schedule_infinite_attempts(self): delay = 2 max_attempts = None crp = ConstantReconnectionPolicy(delay=delay, max_attempts=max_attempts) # this is infinite. we'll just verify one more than default for _, d in zip(range(65), crp.new_schedule()): self.assertEqual(d, delay) class ExponentialReconnectionPolicyTest(unittest.TestCase): def _assert_between(self, value, min, max): self.assertTrue(min <= value <= max) def test_bad_vals(self): self.assertRaises(ValueError, ExponentialReconnectionPolicy, -1, 0) self.assertRaises(ValueError, ExponentialReconnectionPolicy, 0, -1) self.assertRaises(ValueError, ExponentialReconnectionPolicy, 9000, 1) self.assertRaises(ValueError, ExponentialReconnectionPolicy, 1, 2, -1) def test_schedule_no_max(self): base_delay = 2.0 max_delay = 100.0 test_iter = 10000 policy = ExponentialReconnectionPolicy(base_delay=base_delay, max_delay=max_delay, max_attempts=None) sched_slice = list(islice(policy.new_schedule(), 0, test_iter)) self._assert_between(sched_slice[0], base_delay*0.85, base_delay*1.15) self._assert_between(sched_slice[-1], max_delay*0.85, max_delay*1.15) self.assertEqual(len(sched_slice), test_iter) def test_schedule_with_max(self): base_delay = 2.0 max_delay = 100.0 max_attempts = 64 policy = ExponentialReconnectionPolicy(base_delay=base_delay, max_delay=max_delay, max_attempts=max_attempts) schedule = list(policy.new_schedule()) self.assertEqual(len(schedule), max_attempts) for i, delay in enumerate(schedule): if i == 0: self._assert_between(delay, base_delay*0.85, base_delay*1.15) elif i < 6: value = base_delay * (2 ** i) self._assert_between(delay, value*85/100, value*1.15) else: self._assert_between(delay, max_delay*85/100, max_delay*1.15) def test_schedule_exactly_one_attempt(self): base_delay = 2.0 max_delay = 100.0 max_attempts = 1 policy = ExponentialReconnectionPolicy( base_delay=base_delay, max_delay=max_delay, max_attempts=max_attempts ) self.assertEqual(len(list(policy.new_schedule())), 1) def test_schedule_overflow(self): """ Test to verify an OverflowError is handled correctly in the ExponentialReconnectionPolicy @since 3.10 @jira_ticket PYTHON-707 @expected_result all numbers should be less than sys.float_info.max since that's the biggest max we can possibly have as that argument must be a float. Note that is possible for a float to be inf. @test_category policy """ # This should lead to overflow # Note that this may not happen in the fist iterations # as sys.float_info.max * 2 = inf base_delay = sys.float_info.max - 1 max_delay = sys.float_info.max max_attempts = 2**12 policy = ExponentialReconnectionPolicy(base_delay=base_delay, max_delay=max_delay, max_attempts=max_attempts) schedule = list(policy.new_schedule()) for number in schedule: self.assertLessEqual(number, sys.float_info.max) def test_schedule_with_jitter(self): """ Test to verify jitter is added properly and is always between -/+ 15%. @since 3.18 @jira_ticket PYTHON-1065 """ for i in range(100): base_delay = float(randint(2, 5)) max_delay = (base_delay - 1) * 100.0 ep = ExponentialReconnectionPolicy(base_delay, max_delay, max_attempts=64) schedule = ep.new_schedule() for i in range(64): exp_delay = min(base_delay * (2 ** i), max_delay) min_jitter_delay = max(base_delay, exp_delay*85/100) max_jitter_delay = min(max_delay, exp_delay*115/100) delay = next(schedule) self._assert_between(delay, min_jitter_delay, max_jitter_delay) ONE = ConsistencyLevel.ONE class RetryPolicyTest(unittest.TestCase): def test_read_timeout(self): policy = RetryPolicy() # if this is the second or greater attempt, rethrow retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=1, received_responses=2, data_retrieved=True, retry_num=1) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) # if we didn't get enough responses, rethrow retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=2, received_responses=1, data_retrieved=True, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) # if we got enough responses, but also got a data response, rethrow retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=2, received_responses=2, data_retrieved=True, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) # we got enough responses but no data response, so retry retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=2, received_responses=2, data_retrieved=False, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY) self.assertEqual(consistency, ONE) def test_write_timeout(self): policy = RetryPolicy() # if this is the second or greater attempt, rethrow retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=WriteType.SIMPLE, required_responses=1, received_responses=2, retry_num=1) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) # if it's not a BATCH_LOG write, don't retry it retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=WriteType.SIMPLE, required_responses=1, received_responses=2, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) # retry BATCH_LOG writes regardless of received responses retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=WriteType.BATCH_LOG, required_responses=10000, received_responses=1, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY) self.assertEqual(consistency, ONE) def test_unavailable(self): """ Use the same tests for test_write_timeout, but ensure they only RETHROW """ policy = RetryPolicy() retry, consistency = policy.on_unavailable( query=None, consistency=ONE, required_replicas=1, alive_replicas=2, retry_num=1) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) retry, consistency = policy.on_unavailable( query=None, consistency=ONE, required_replicas=1, alive_replicas=2, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY_NEXT_HOST) self.assertEqual(consistency, None) retry, consistency = policy.on_unavailable( query=None, consistency=ONE, required_replicas=10000, alive_replicas=1, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY_NEXT_HOST) self.assertEqual(consistency, None) class FallthroughRetryPolicyTest(unittest.TestCase): """ Use the same tests for test_write_timeout, but ensure they only RETHROW """ def test_read_timeout(self): policy = FallthroughRetryPolicy() retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=1, received_responses=2, data_retrieved=True, retry_num=1) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=2, received_responses=1, data_retrieved=True, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=2, received_responses=2, data_retrieved=True, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=2, received_responses=2, data_retrieved=False, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) def test_write_timeout(self): policy = FallthroughRetryPolicy() retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=WriteType.SIMPLE, required_responses=1, received_responses=2, retry_num=1) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=WriteType.SIMPLE, required_responses=1, received_responses=2, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=WriteType.BATCH_LOG, required_responses=10000, received_responses=1, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) def test_unavailable(self): policy = FallthroughRetryPolicy() retry, consistency = policy.on_unavailable( query=None, consistency=ONE, required_replicas=1, alive_replicas=2, retry_num=1) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) retry, consistency = policy.on_unavailable( query=None, consistency=ONE, required_replicas=1, alive_replicas=2, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) retry, consistency = policy.on_unavailable( query=None, consistency=ONE, required_replicas=10000, alive_replicas=1, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) class DowngradingConsistencyRetryPolicyTest(unittest.TestCase): def test_read_timeout(self): policy = DowngradingConsistencyRetryPolicy() # if this is the second or greater attempt, rethrow retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=1, received_responses=2, data_retrieved=True, retry_num=1) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) # if we didn't get enough responses, retry at a lower consistency retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=4, received_responses=3, data_retrieved=True, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY) self.assertEqual(consistency, ConsistencyLevel.THREE) # if we didn't get enough responses, retry at a lower consistency retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=3, received_responses=2, data_retrieved=True, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY) self.assertEqual(consistency, ConsistencyLevel.TWO) # retry consistency level goes down based on the # of recv'd responses retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=3, received_responses=1, data_retrieved=True, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY) self.assertEqual(consistency, ConsistencyLevel.ONE) # if we got no responses, rethrow retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=3, received_responses=0, data_retrieved=True, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) # if we got enough response but no data, retry retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=3, received_responses=3, data_retrieved=False, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY) self.assertEqual(consistency, ONE) # if we got enough responses, but also got a data response, rethrow retry, consistency = policy.on_read_timeout( query=None, consistency=ONE, required_responses=2, received_responses=2, data_retrieved=True, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) def test_write_timeout(self): policy = DowngradingConsistencyRetryPolicy() # if this is the second or greater attempt, rethrow retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=WriteType.SIMPLE, required_responses=1, received_responses=2, retry_num=1) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) for write_type in (WriteType.SIMPLE, WriteType.BATCH, WriteType.COUNTER): # ignore failures if at least one response (replica persisted) retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=write_type, required_responses=1, received_responses=2, retry_num=0) self.assertEqual(retry, RetryPolicy.IGNORE) # retrhow if we can't be sure we have a replica retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=write_type, required_responses=1, received_responses=0, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) # downgrade consistency level on unlogged batch writes retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=WriteType.UNLOGGED_BATCH, required_responses=3, received_responses=1, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY) self.assertEqual(consistency, ConsistencyLevel.ONE) # retry batch log writes at the same consistency level retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=WriteType.BATCH_LOG, required_responses=3, received_responses=1, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY) self.assertEqual(consistency, ONE) # timeout on an unknown write_type retry, consistency = policy.on_write_timeout( query=None, consistency=ONE, write_type=None, required_responses=1, received_responses=2, retry_num=0) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) def test_unavailable(self): policy = DowngradingConsistencyRetryPolicy() # if this is the second or greater attempt, rethrow retry, consistency = policy.on_unavailable( query=None, consistency=ONE, required_replicas=3, alive_replicas=1, retry_num=1) self.assertEqual(retry, RetryPolicy.RETHROW) self.assertEqual(consistency, None) # downgrade consistency on unavailable exceptions retry, consistency = policy.on_unavailable( query=None, consistency=ONE, required_replicas=3, alive_replicas=1, retry_num=0) self.assertEqual(retry, RetryPolicy.RETRY) self.assertEqual(consistency, ConsistencyLevel.ONE) class WhiteListRoundRobinPolicyTest(unittest.TestCase): def test_hosts_with_hostname(self): hosts = ['localhost'] policy = WhiteListRoundRobinPolicy(hosts) host = Host(DefaultEndPoint("127.0.0.1"), SimpleConvictionPolicy) policy.populate(None, [host]) qplan = list(policy.make_query_plan()) self.assertEqual(sorted(qplan), [host]) self.assertEqual(policy.distance(host), HostDistance.LOCAL) class AddressTranslatorTest(unittest.TestCase): def test_identity_translator(self): IdentityTranslator() @patch('socket.getfqdn', return_value='localhost') def test_ec2_multi_region_translator(self, *_): ec2t = EC2MultiRegionTranslator() addr = '127.0.0.1' translated = ec2t.translate(addr) self.assertIsNot(translated, addr) # verifies that the resolver path is followed self.assertEqual(translated, addr) # and that it resolves to the same address class HostFilterPolicyInitTest(unittest.TestCase): def setUp(self): self.child_policy, self.predicate = (Mock(name='child_policy'), Mock(name='predicate')) def _check_init(self, hfp): self.assertIs(hfp._child_policy, self.child_policy) self.assertIsInstance(hfp._hosts_lock, LockType) # we can't use a simple assertIs because we wrap the function arg0, arg1 = Mock(name='arg0'), Mock(name='arg1') hfp.predicate(arg0) hfp.predicate(arg1) self.predicate.assert_has_calls([call(arg0), call(arg1)]) def test_init_arg_order(self): self._check_init(HostFilterPolicy(self.child_policy, self.predicate)) def test_init_kwargs(self): self._check_init(HostFilterPolicy( predicate=self.predicate, child_policy=self.child_policy )) def test_immutable_predicate(self): if sys.version_info >= (3, 11): expected_message_regex = "has no setter" else: expected_message_regex = "can't set attribute" hfp = HostFilterPolicy(child_policy=Mock(name='child_policy'), predicate=Mock(name='predicate')) with self.assertRaisesRegex(AttributeError, expected_message_regex): hfp.predicate = object() class HostFilterPolicyDeferralTest(unittest.TestCase): def setUp(self): self.passthrough_hfp = HostFilterPolicy( child_policy=Mock(name='child_policy'), predicate=Mock(name='passthrough_predicate', return_value=True) ) self.filterall_hfp = HostFilterPolicy( child_policy=Mock(name='child_policy'), predicate=Mock(name='filterall_predicate', return_value=False) ) def _check_host_triggered_method(self, policy, name): arg, kwarg = Mock(name='arg'), Mock(name='kwarg') method, child_policy_method = (getattr(policy, name), getattr(policy._child_policy, name)) result = method(arg, kw=kwarg) # method calls the child policy's method... child_policy_method.assert_called_once_with(arg, kw=kwarg) # and returns its return value self.assertIs(result, child_policy_method.return_value) def test_defer_on_up_to_child_policy(self): self._check_host_triggered_method(self.passthrough_hfp, 'on_up') def test_defer_on_down_to_child_policy(self): self._check_host_triggered_method(self.passthrough_hfp, 'on_down') def test_defer_on_add_to_child_policy(self): self._check_host_triggered_method(self.passthrough_hfp, 'on_add') def test_defer_on_remove_to_child_policy(self): self._check_host_triggered_method(self.passthrough_hfp, 'on_remove') def test_filtered_host_on_up_doesnt_call_child_policy(self): self._check_host_triggered_method(self.filterall_hfp, 'on_up') def test_filtered_host_on_down_doesnt_call_child_policy(self): self._check_host_triggered_method(self.filterall_hfp, 'on_down') def test_filtered_host_on_add_doesnt_call_child_policy(self): self._check_host_triggered_method(self.filterall_hfp, 'on_add') def test_filtered_host_on_remove_doesnt_call_child_policy(self): self._check_host_triggered_method(self.filterall_hfp, 'on_remove') def _check_check_supported_deferral(self, policy): policy.check_supported() policy._child_policy.check_supported.assert_called_once() def test_check_supported_defers_to_child(self): self._check_check_supported_deferral(self.passthrough_hfp) def test_check_supported_defers_to_child_when_predicate_filtered(self): self._check_check_supported_deferral(self.filterall_hfp) class HostFilterPolicyDistanceTest(unittest.TestCase): def setUp(self): self.hfp = HostFilterPolicy( child_policy=Mock(name='child_policy', distance=Mock(name='distance')), predicate=lambda host: host.address == 'acceptme' ) self.ignored_host = Host(DefaultEndPoint('ignoreme'), conviction_policy_factory=Mock()) self.accepted_host = Host(DefaultEndPoint('acceptme'), conviction_policy_factory=Mock()) def test_ignored_with_filter(self): self.assertEqual(self.hfp.distance(self.ignored_host), HostDistance.IGNORED) self.assertNotEqual(self.hfp.distance(self.accepted_host), HostDistance.IGNORED) def test_accepted_filter_defers_to_child_policy(self): self.hfp._child_policy.distance.side_effect = distances = Mock(), Mock() # getting the distance for an ignored host shouldn't affect subsequent results self.hfp.distance(self.ignored_host) # first call of _child_policy with count() side effect self.assertEqual(self.hfp.distance(self.accepted_host), distances[0]) # second call of _child_policy with count() side effect self.assertEqual(self.hfp.distance(self.accepted_host), distances[1]) class HostFilterPolicyPopulateTest(unittest.TestCase): def test_populate_deferred_to_child(self): hfp = HostFilterPolicy( child_policy=Mock(name='child_policy'), predicate=lambda host: True ) mock_cluster, hosts = (Mock(name='cluster'), ['host1', 'host2', 'host3']) hfp.populate(mock_cluster, hosts) hfp._child_policy.populate.assert_called_once_with( cluster=mock_cluster, hosts=hosts ) def test_child_is_populated_with_filtered_hosts(self): hfp = HostFilterPolicy( child_policy=Mock(name='child_policy'), predicate=lambda host: False ) mock_cluster, hosts = (Mock(name='cluster'), ['acceptme0', 'acceptme1']) hfp.populate(mock_cluster, hosts) hfp._child_policy.populate.assert_called_once() self.assertEqual( hfp._child_policy.populate.call_args[1]['hosts'], ['acceptme0', 'acceptme1'] ) class HostFilterPolicyQueryPlanTest(unittest.TestCase): def test_query_plan_deferred_to_child(self): child_policy = Mock( name='child_policy', make_query_plan=Mock( return_value=[object(), object(), object()] ) ) hfp = HostFilterPolicy( child_policy=child_policy, predicate=lambda host: True ) working_keyspace, query = (Mock(name='working_keyspace'), Mock(name='query')) qp = list(hfp.make_query_plan(working_keyspace=working_keyspace, query=query)) hfp._child_policy.make_query_plan.assert_called_once_with( working_keyspace=working_keyspace, query=query ) self.assertEqual(qp, hfp._child_policy.make_query_plan.return_value) def test_wrap_token_aware(self): cluster = Mock(spec=Cluster) hosts = [Host(DefaultEndPoint("127.0.0.{}".format(i)), SimpleConvictionPolicy) for i in range(1, 6)] for host in hosts: host.set_up() def get_replicas(keyspace, packed_key): return hosts[:2] cluster.metadata.get_replicas.side_effect = get_replicas child_policy = TokenAwarePolicy(RoundRobinPolicy()) hfp = HostFilterPolicy( child_policy=child_policy, predicate=lambda host: host.address != "127.0.0.1" and host.address != "127.0.0.4" ) hfp.populate(cluster, hosts) # We don't allow randomness for ordering the replicas in RoundRobin hfp._child_policy._child_policy._position = 0 mocked_query = Mock() query_plan = hfp.make_query_plan("keyspace", mocked_query) # First the not filtered replica, and then the rest of the allowed hosts ordered query_plan = list(query_plan) self.assertEqual(query_plan[0], Host(DefaultEndPoint("127.0.0.2"), SimpleConvictionPolicy)) self.assertEqual(set(query_plan[1:]),{Host(DefaultEndPoint("127.0.0.3"), SimpleConvictionPolicy), Host(DefaultEndPoint("127.0.0.5"), SimpleConvictionPolicy)}) def test_create_whitelist(self): cluster = Mock(spec=Cluster) hosts = [Host(DefaultEndPoint("127.0.0.{}".format(i)), SimpleConvictionPolicy) for i in range(1, 6)] for host in hosts: host.set_up() child_policy = RoundRobinPolicy() hfp = HostFilterPolicy( child_policy=child_policy, predicate=lambda host: host.address == "127.0.0.1" or host.address == "127.0.0.4" ) hfp.populate(cluster, hosts) # We don't allow randomness for ordering the replicas in RoundRobin hfp._child_policy._position = 0 mocked_query = Mock() query_plan = hfp.make_query_plan("keyspace", mocked_query) # Only the filtered replicas should be allowed self.assertEqual(set(query_plan), {Host(DefaultEndPoint("127.0.0.1"), SimpleConvictionPolicy), Host(DefaultEndPoint("127.0.0.4"), SimpleConvictionPolicy)})