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]))