def _assert_shuffle(self, patched_shuffle, keyspace, routing_key): hosts = [Host(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)
def test_wrap_round_robin(self): cluster = Mock(spec=Cluster) cluster.metadata = Mock(spec=Metadata) hosts = [Host(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_statement_keyspace(self): hosts = [Host(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_status_updates(self): """ Same test as DCAwareRoundRobinPolicyTest.test_status_updates() """ hosts = [Host(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(4, SimpleConvictionPolicy) new_local_host.set_location_info("dc1", "rack1") policy.on_up(new_local_host) new_remote_host = Host(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_wrap_dc_aware(self): cluster = Mock(spec=Cluster) cluster.metadata = Mock(spec=Metadata) hosts = [Host(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))
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("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("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("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]))