async def test_close_only():
pool1 = create_pool_mock()
pool2 = create_pool_mock()
pool3 = create_pool_mock()
mocked_create_pool = mock.AsyncMock(side_effect=[pool1, pool2, pool3])
pooler = Pooler(mocked_create_pool)
addr1 = Address("h1", 1)
addr2 = Address("h2", 2)
result_pool1 = await pooler.ensure_pool(addr1)
await pooler.ensure_pool(addr2)
result_pool2 = await pooler.ensure_pool(addr2)
assert result_pool1 is pool1
assert result_pool2 is pool2
assert mocked_create_pool.call_count == 2
await pooler.close_only([addr2])
pool2.close.assert_called_once_with()
pool2.wait_closed.assert_called_once_with()
result_pool3 = await pooler.ensure_pool(addr2)
assert result_pool3 is pool3
result_pool1 = await pooler.ensure_pool(addr1)
assert result_pool1 is pool1
def test_state__with_fail_state():
state = create_cluster_state(
SLOTS,
{
"cluster_state": "fail",
"cluster_current_epoch": "1",
"cluster_slots_assigned": "16384",
},
Address("172.17.0.1", 6379),
)
assert state.state is NodeClusterState.FAIL
state = create_cluster_state(
SLOTS,
{
"cluster_state": "foobar",
"cluster_current_epoch": "1",
"cluster_slots_assigned": "16384",
},
Address("172.17.0.1", 6379),
)
assert state.state is NodeClusterState.UNKNOWN
assert state.current_epoch == 1
async def test_add_pubsub_channel():
pooler = Pooler(mock.AsyncMock(return_value=create_pool_mock()), reap_frequency=-1)
addr1 = Address("h1", 1234)
addr2 = Address("h2", 1234)
pooler._pubsub_addrs[addr1] = set()
pooler._pubsub_addrs[addr2] = set()
result1 = pooler.add_pubsub_channel(addr1, b"ch1", False)
result2 = pooler.add_pubsub_channel(addr1, b"ch1", True)
result3 = pooler.add_pubsub_channel(addr1, b"ch2", False)
result4 = pooler.add_pubsub_channel(addr2, b"ch3", False)
result5 = pooler.add_pubsub_channel(addr1, b"ch3", False)
assert result1 is True
assert result2 is True
assert result3 is True
assert result4 is True
assert result5 is False
assert len(pooler._pubsub_addrs[addr1]) == 3
assert len(pooler._pubsub_addrs[addr2]) == 1
assert len(pooler._pubsub_channels) == 4
collected_channels = [(ch.name, ch.is_pattern) for ch in pooler._pubsub_channels]
assert (b"ch1", False) in collected_channels
assert (b"ch1", True) in collected_channels
assert (b"ch2", False) in collected_channels
assert (b"ch3", False) in collected_channels
async def test_all_masters__with_attempt_timeout(mocker, event_loop):
cl = Cluster(["addr1"], attempt_timeout=0.001)
mocked_pooler = mocker.patch.object(cl, "_pooler", new=get_pooler_mock())
pool = mocked_pooler._pool
execute1_fut = event_loop.create_future()
execute1_fut.set_result(None)
execute2_fut = event_loop.create_future()
execute_futs_iter = iter([execute1_fut, execute2_fut])
async def execute_side_effect(*args, **kwargs):
return await next(execute_futs_iter)
pool.execute.side_effect = execute_side_effect
mocked_manager = mocker.patch.object(cl,
"_manager",
new=get_manager_mock())
state = mocked_manager.get_state.return_value
state._data.masters = [
mock.NonCallableMock(addr=Address("1.2.3.4", 6666)),
mock.NonCallableMock(addr=Address("1.2.3.4", 9999)),
]
mocked_execute_retry_slowdown = mocker.patch.object(
cl, "_execute_retry_slowdown", new=create_async_mock())
with pytest.raises(asyncio.TimeoutError):
await cl.all_masters()
assert mocked_manager.get_state.call_count == 1
assert mocked_pooler.ensure_pool.call_count == 2
assert mocked_manager.require_reload_state.call_count == 1
mocked_execute_retry_slowdown.assert_not_called()
async def test_remove_pubsub_channel():
pooler = Pooler(mock.AsyncMock(), reap_frequency=-1)
addr1 = Address("h1", 1234)
addr2 = Address("h2", 1234)
pooler._pubsub_addrs[addr1] = set()
pooler._pubsub_addrs[addr2] = set()
pooler.add_pubsub_channel(addr1, b"ch1", False)
pooler.add_pubsub_channel(addr1, b"ch2", False)
pooler.add_pubsub_channel(addr2, b"ch3", True)
result1 = pooler.remove_pubsub_channel(b"ch1", False)
result2 = pooler.remove_pubsub_channel(b"ch1", True)
result3 = pooler.remove_pubsub_channel(b"ch2", False)
result4 = pooler.remove_pubsub_channel(b"ch3", True)
result5 = pooler.remove_pubsub_channel(b"ch3", True)
assert result1 is True
assert result2 is False
assert result3 is True
assert result4 is True
assert result5 is False
assert len(pooler._pubsub_addrs[addr1]) == 0
assert len(pooler._pubsub_addrs[addr2]) == 0
assert len(pooler._pubsub_channels) == 0
async def test_reap_pools(mocker):
addrs_pools = [
(Address("h1", 1), create_pool_mock()),
(Address("h2", 2), create_pool_mock()),
]
addrs = [p[0] for p in addrs_pools]
pools = [p[1] for p in addrs_pools]
mocked_create_pool = mock.AsyncMock(side_effect=pools)
pooler = Pooler(mocked_create_pool, reap_frequency=-1)
# create pools
await pooler.ensure_pool(addrs[0])
await pooler.ensure_pool(addrs[1])
# try to reap pools
reaped = await pooler._reap_pools()
assert len(reaped) == 0
# touch only one pool
await pooler.ensure_pool(addrs[1])
reaped = await pooler._reap_pools()
assert len(reaped) == 1
assert reaped[0] is pools[0]
assert len(pooler._nodes) == 1
reaped = await pooler._reap_pools()
assert len(reaped) == 1
assert reaped[0] is pools[1]
assert len(pooler._nodes) == 0
async def test_create_cluster__defaults(mocker, cluster_fix):
cluster = cluster_fix()
mocked_cluster_cls = mocker.patch(
create_cluster.__module__ + ".Cluster",
return_value=cluster,
)
result = await create_cluster([
"redis://redis1?db=1",
("redis2", 6380),
])
assert result is cluster
mocked_cluster_cls.assert_called_once()
assert len(mocked_cluster_cls.call_args[0][0]) == 2
assert Address("redis1", 6379) in mocked_cluster_cls.call_args[0][0]
assert Address("redis2", 6380) in mocked_cluster_cls.call_args[0][0]
assert mocked_cluster_cls.call_args == mock.call(
mock.ANY,
retry_min_delay=None,
retry_max_delay=None,
max_attempts=None,
state_reload_interval=None,
follow_cluster=None,
idle_connection_timeout=None,
password=None,
encoding=None,
pool_minsize=None,
pool_maxsize=None,
connect_timeout=None,
attempt_timeout=None,
)
cluster._init.assert_called_once()
cluster.wait_closed.assert_not_called()
async def test_execute__with_attempt_timeout__idempotent__success(mocker):
cl = Cluster(["addr1"], max_attempts=4)
mocked_pooler = mocker.patch.object(cl, "_pooler", new=get_pooler_mock())
pool = mocked_pooler._pool
pool.execute.side_effect = [
asyncio.TimeoutError(),
asyncio.TimeoutError(),
asyncio.TimeoutError(),
b"VALUE",
]
mocked_manager = mocker.patch.object(cl,
"_manager",
new=get_manager_mock())
state = mocked_manager.get_state.return_value
state.slot_master.return_value.addr = Address("1.2.3.4", 9999)
state.random_node.return_value.addr = Address("6.6.6.6", 9999)
mocked_execute_retry_slowdown = mocker.patch.object(
cl, "_execute_retry_slowdown", new=create_async_mock())
result = await cl.execute("get", "key")
assert result == b"VALUE"
assert mocked_manager.get_state.call_count == 4
assert mocked_pooler.ensure_pool.call_count == 4
assert mocked_manager.require_reload_state.call_count == 3
assert mocked_execute_retry_slowdown.call_count == 3
def _make_execute_props(
self,
state: ClusterState,
ctx: ExecuteContext,
fail_props: ExecuteFailProps = None,
) -> ExecuteProps:
exec_props = ExecuteProps()
node_addr: Address
if fail_props:
# reraise exception for simplify classification
# instead of many isinstance conditions
try:
raise fail_props.error
except self._connection_errors:
if ctx.attempt <= 2 and ctx.slot is not None:
replica = state.random_slot_replica(ctx.slot)
if replica is not None:
node_addr = replica.addr
else:
node_addr = state.random_node().addr
else:
node_addr = state.random_node().addr
except MovedError as e:
node_addr = Address(e.info.host, e.info.port)
except AskError as e:
node_addr = Address(e.info.host, e.info.port)
exec_props.asking = e.info.ask
except (ClusterDownError, TryAgainError, LoadingError, ProtocolError):
node_addr = state.random_node().addr
except Exception as e:
# usualy never be done here
logger.exception("Uncaught exception on execute: %r", e)
raise
logger.info("New node to execute: %s", node_addr)
else:
if ctx.slot is not None:
try:
node = state.slot_master(ctx.slot)
except UncoveredSlotError:
logger.warning("No node found by slot %d", ctx.slot)
# probably cluster is corrupted and
# we need try to recover cluster state
exec_props.reload_state_required = True
node = state.random_master()
node_addr = node.addr
else:
node_addr = state.random_master().addr
logger.debug("Defined node to command: %s", node_addr)
exec_props.node_addr = node_addr
return exec_props
async def test_execute__retryable_errors(mocker, error, retry_node_addr,
require_reload_state):
cl = Cluster(["addr1"])
mocked_pooler = mocker.patch.object(cl, "_pooler", new=get_pooler_mock())
pool = mocked_pooler._pool
conn = mocked_pooler._conn
pool_execute_count = 0
def pool_execute_se(*args, **kwargs):
nonlocal pool_execute_count
if args[0] == b"ASKING":
return b"OK"
pool_execute_count += 1
if pool_execute_count == 1:
raise error
return pool.execute.return_value
pool.execute.side_effect = pool_execute_se
conn.execute.side_effect = pool_execute_se
mocked_manager = mocker.patch.object(cl,
"_manager",
new=get_manager_mock())
state = mocked_manager.get_state.return_value
state.slot_master.return_value.addr = Address("master", 6379)
state.random_node.return_value.addr = Address("random", 6379)
state.random_slot_replica.return_value.addr = Address(
"random_replica", 6379)
mocked_execute_retry_slowdown = mocker.patch.object(
cl, "_execute_retry_slowdown", new=create_async_mock())
result = await cl.execute("get", "key")
assert mocked_manager.get_state.call_count == 2
assert mocked_pooler.ensure_pool.call_count == 2
assert mocked_manager.require_reload_state.call_count == require_reload_state
assert result is pool.execute.return_value
assert mocked_pooler.ensure_pool.call_args_list[0] == mock.call(
Address("master", 6379))
assert mocked_pooler.ensure_pool.call_args_list[1] == mock.call(
retry_node_addr)
# cl.determine_slot(b'key') == 12539
state.slot_master.assert_called_once_with(12539)
mocked_execute_retry_slowdown.assert_called_once_with(1, 10)
async def test_ensure_pool__identical_address():
mocked_create_pool = mock.AsyncMock(
return_value=create_pool_mock(),
)
pooler = Pooler(mocked_create_pool)
result = await pooler.ensure_pool(Address("localhost", 1234))
assert result is mocked_create_pool.return_value
mocked_create_pool.assert_called_once_with(("localhost", 1234))
result2 = await pooler.ensure_pool(Address("localhost", 1234))
assert result2 is result
assert mocked_create_pool.call_count == 1
async def test_ensure_pool__only_one(event_loop):
pools = {
("h1", 1): create_pool_mock(),
("h2", 2): create_pool_mock(),
}
pool_creation_fut = event_loop.create_future()
async def create_pool_se(addr):
nonlocal pool_creation_fut
await pool_creation_fut
return pools[addr]
mocked_create_pool = mock.AsyncMock(side_effect=create_pool_se)
pooler = Pooler(mocked_create_pool)
tasks = []
for i in range(10):
for addr in pools.keys():
task = event_loop.create_task(pooler.ensure_pool(Address(addr[0], addr[1])))
tasks.append(task)
pool_creation_fut.set_result(None)
results = await asyncio.gather(*tasks)
assert len(results) == 20
assert len([r for r in results if r is pools[("h1", 1)]]) == 10
assert len([r for r in results if r is pools[("h2", 2)]]) == 10
assert mocked_create_pool.call_count == 2
async def test_add_pubsub_channel__no_addr():
pooler = Pooler(mock.AsyncMock(), reap_frequency=-1)
addr = Address("h1", 1234)
result = pooler.add_pubsub_channel(addr, b"channel", False)
assert result is False
async def test_get_pubsub_addr():
pooler = Pooler(mock.AsyncMock(), reap_frequency=-1)
addr1 = Address("h1", 1234)
addr2 = Address("h2", 1234)
pooler._pubsub_addrs[addr1] = set()
pooler._pubsub_addrs[addr2] = set()
pooler.add_pubsub_channel(addr1, b"ch1", False)
pooler.add_pubsub_channel(addr2, b"ch2", True)
result1 = pooler.get_pubsub_addr(b"ch1", False)
result2 = pooler.get_pubsub_addr(b"ch1", True)
result3 = pooler.get_pubsub_addr(b"ch2", False)
result4 = pooler.get_pubsub_addr(b"ch2", True)
assert result1 == addr1
assert result2 is None
assert result3 is None
assert result4 == addr2
async def test_ensure_pool__create_pubsub_addr_set():
addr1 = Address("h1", 1234)
addr2 = Address("h2", 1234)
pooler = Pooler(mock.AsyncMock(return_value=create_pool_mock()))
assert len(pooler._pubsub_addrs) == 0
await pooler.ensure_pool(addr1)
await pooler.ensure_pool(addr2)
await pooler.ensure_pool(addr2)
assert len(pooler._pubsub_addrs) == 2
assert addr1 in pooler._pubsub_addrs
assert addr2 in pooler._pubsub_addrs
assert len(pooler._pubsub_addrs[addr1]) == 0
pooler._pubsub_addrs[addr1].add(object())
await pooler.ensure_pool(addr1)
assert len(pooler._pubsub_addrs[addr1]) == 1
async def test_ensure_pool__multiple():
pools = [object(), object(), object()]
mocked_create_pool = mock.AsyncMock(side_effect=pools)
pooler = Pooler(mocked_create_pool)
result1 = await pooler.ensure_pool(Address("localhost", 1234))
result2 = await pooler.ensure_pool(Address("localhost", 4321))
result3 = await pooler.ensure_pool(Address("127.0.0.1", 1234))
assert result1 is pools[0]
assert result2 is pools[1]
assert result3 is pools[2]
assert mocked_create_pool.call_count == 3
mocked_create_pool.assert_has_calls(
[
mock.call(("localhost", 1234)),
mock.call(("localhost", 4321)),
mock.call(("127.0.0.1", 1234)),
]
)
def test_create_cluster_state():
state = create_cluster_state(
SLOTS,
INFO,
Address("172.17.0.1", 6379),
)
addrs = sorted([Address("172.17.0.2", port) for port in range(7000, 7005)])
masters_addrs = sorted([
Address("172.17.0.2", 7000),
Address("172.17.0.2", 7001),
Address("172.17.0.2", 7002),
])
replicas_addrs = sorted([
Address("172.17.0.2", 7003),
Address("172.17.0.2", 7004),
])
slot_ranges = sorted([
(0, 5460),
(9995, 9995),
(12182, 12182),
(5461, 9994),
(9996, 10922),
(10923, 12181),
(12183, 16383),
])
state_data = state._data
assert len(state_data.nodes) == 5
assert len(state_data.addrs) == 5
assert len(set(state_data.addrs)) == 5
assert len(state_data.masters) == 3
assert len(state_data.replicas) == 3
assert sum(len(rs) for rs in state_data.replicas.values()) == 2
assert len(state_data.slots) == 7
assert sorted(state_data.addrs) == addrs
assert sorted(state_data.nodes.keys()) == addrs
assert get_nodes_addr(state_data.masters) == masters_addrs
assert get_nodes_addr([
r for rs in state_data.replicas.values() for r in rs
]) == replicas_addrs
assert get_slots_ranges(state_data.slots) == slot_ranges
assert state.state is NodeClusterState.OK
assert state.state_from == Address("172.17.0.1", 6379)
assert isinstance(state._data.created_at, datetime.datetime)
assert state._data.created_at_local > 0
assert state.current_epoch == 1
state.repr_stats()
str(state)
async def test_reap_pools__cleanup_channels():
pooler = Pooler(mock.AsyncMock(), reap_frequency=-1)
addr1 = Address("h1", 1)
addr2 = Address("h2", 2)
# create pools
await pooler.ensure_pool(addr1)
await pooler.ensure_pool(addr2)
pooler.add_pubsub_channel(addr1, b"ch1", False)
pooler.add_pubsub_channel(addr2, b"ch2", False)
# try to reap pools
reaped = await pooler._reap_pools()
assert len(reaped) == 0
reaped = await pooler._reap_pools()
assert len(reaped) == 2
assert len(pooler._pubsub_addrs) == 0
assert len(pooler._pubsub_channels) == 0
async def test_close__with_pools(mocker):
addrs_pools = [
(Address("h1", 1), create_pool_mock()),
(Address("h2", 2), create_pool_mock()),
]
addrs = [p[0] for p in addrs_pools]
pools = [p[1] for p in addrs_pools]
mocked_create_pool = mock.AsyncMock(side_effect=pools)
pooler = Pooler(mocked_create_pool)
result1 = await pooler.ensure_pool(addrs[0])
result2 = await pooler.ensure_pool(addrs[1])
assert len(pooler._nodes) == 2
await pooler.close()
assert len(pooler._nodes) == 0
assert pooler.closed is True
result1.close.assert_called_once()
result2.close.assert_called_once()
result1.wait_closed.assert_called_once()
result2.wait_closed.assert_called_once()
async def test_create_cluster__parametrized(mocker, cluster_fix):
cluster = cluster_fix()
mocked_cluster_cls = mocker.patch(
create_cluster.__module__ + ".Cluster",
return_value=cluster,
)
result = await create_cluster(
("redis://redis1?db=1", ),
retry_min_delay=1.2,
retry_max_delay=2.3,
max_attempts=8,
attempt_timeout=5.5,
state_reload_frequency=7.8,
state_reload_interval=8.7,
follow_cluster=True,
idle_connection_timeout=1.1,
password="******",
encoding="utf-8",
pool_minsize=1,
pool_maxsize=5,
connect_timeout=3.5,
)
assert result is cluster
mocked_cluster_cls.assert_called_once()
assert mocked_cluster_cls.call_args == mock.call(
[
Address("redis1", 6379),
],
retry_min_delay=1.2,
retry_max_delay=2.3,
max_attempts=8,
attempt_timeout=5.5,
state_reload_interval=8.7,
follow_cluster=True,
idle_connection_timeout=1.1,
password="******",
encoding="utf-8",
pool_minsize=1,
pool_maxsize=5,
connect_timeout=3.5,
)
async def test_ensure_pool__error():
pools = [RuntimeError(), object()]
mocked_create_pool = mock.AsyncMock(side_effect=pools)
pooler = Pooler(mocked_create_pool)
addr = Address("localhost", 1234)
with pytest.raises(RuntimeError):
await pooler.ensure_pool(addr)
result = await pooler.ensure_pool(addr)
assert result is pools[1]
assert mocked_create_pool.call_count == 2
mocked_create_pool.assert_has_calls(
[
mock.call(("localhost", 1234)),
mock.call(("localhost", 1234)),
]
)
async def test_create_pool_by_addr__bad_addr(mocker):
pool_cls = mock.Mock()
cl = Cluster(
["addr1"],
password="******",
encoding="utf-8",
pool_minsize=5,
pool_maxsize=15,
connect_timeout=1.2,
pool_cls=pool_cls,
)
mocked_manager = mocker.patch.object(cl,
"_manager",
new=get_manager_mock())
state = mocked_manager.get_state.return_value
state.has_addr.return_value = False
addr = Address("addr4", 8004)
with pytest.raises(ValueError):
await cl.create_pool_by_addr(addr, maxsize=42)
state.has_addr.assert_called_once_with(addr)
async def test_keys_master__with_attempt_timeout(mocker):
cl = Cluster(["addr1"], max_attempts=2)
mocked_pooler = mocker.patch.object(cl, "_pooler", new=get_pooler_mock())
pool = mocked_pooler._pool
pool.execute.side_effect = asyncio.TimeoutError()
mocked_manager = mocker.patch.object(cl,
"_manager",
new=get_manager_mock())
state = mocked_manager.get_state.return_value
state.slot_master.return_value.addr = Address("1.2.3.4", 9999)
mocked_execute_retry_slowdown = mocker.patch.object(
cl, "_execute_retry_slowdown", new=create_async_mock())
with pytest.raises(asyncio.TimeoutError):
await cl.keys_master("foo{key}", "bar{key}")
assert mocked_manager.get_state.call_count == 2
assert mocked_pooler.ensure_pool.call_count == 2
assert mocked_manager.require_reload_state.call_count == 2
assert mocked_execute_retry_slowdown.call_count == 1
async def test_create_pool_by_addr(mocker):
mocked_create_pool = mocker.patch(Cluster.__module__ + ".create_pool",
new=create_async_mock())
pool_cls = mock.Mock()
cl = Cluster(
["addr1"],
password="******",
encoding="utf-8",
pool_minsize=5,
pool_maxsize=15,
connect_timeout=1.2,
pool_cls=pool_cls,
)
mocked_manager = mocker.patch.object(cl,
"_manager",
new=get_manager_mock())
state = mocked_manager.get_state.return_value
state._data.masters = [
mock.NonCallableMock(name="master1", addr=("addr1", 8001)),
mock.NonCallableMock(name="master2", addr=("addr2", 8002)),
mock.NonCallableMock(name="master3", addr=("addr3", 8003)),
]
addr = Address("addr2", 8002)
redis = await cl.create_pool_by_addr(addr, maxsize=42)
assert redis.connection is mocked_create_pool.return_value
mocked_create_pool.assert_called_once_with(
("addr2", 8002),
pool_cls=pool_cls,
password="******",
encoding="utf-8",
minsize=5,
maxsize=42,
create_connection_timeout=1.2,
)
async def test_execute__with_attempt_timeout__non_idempotent(
mocker, event_loop):
cl = Cluster(["addr1"], attempt_timeout=0.001)
mocked_pooler = mocker.patch.object(cl, "_pooler", new=get_pooler_mock())
pool = mocked_pooler._pool
pool.execute = mock.Mock(return_value=event_loop.create_future())
mocked_manager = mocker.patch.object(cl,
"_manager",
new=get_manager_mock())
state = mocked_manager.get_state.return_value
state.slot_master.return_value.addr = Address("1.2.3.4", 9999)
mocked_execute_retry_slowdown = mocker.patch.object(
cl, "_execute_retry_slowdown", new=create_async_mock())
with pytest.raises(asyncio.TimeoutError):
await cl.execute("set", "key", "value")
assert mocked_manager.get_state.call_count == 1
assert mocked_pooler.ensure_pool.call_count == 1
assert mocked_manager.require_reload_state.call_count == 1
mocked_execute_retry_slowdown.assert_not_called()
async def test_load_state(mocker, pooler_mock):
pooler = pooler_mock()
mocker.patch(ClusterManager.__module__ + ".random.sample",
return_value=["addr1", "addr2"])
manager = ClusterManager(["addr1", "addr2"], pooler)
mocked_state = mock.NonCallableMock()
mocked_state._data.masters = [
ClusterNode(Address("master1", 6666), "master1_id"),
ClusterNode(Address("master2", 7777), "master2_id"),
]
mocked_state._data.addrs = [object(), object()]
mocked_state.random_master.return_value = ClusterNode(
Address("random_master", 5555), "random_master_id")
mocked_fetch_state = mocker.patch.object(
manager,
"_fetch_state",
new=mock.AsyncMock(return_value=mocked_state, ),
)
mocked_create_registry = mocker.patch(manager.__module__ +
".create_registry")
result = await manager._load_state(1)
assert result is mocked_state
assert result is manager._state
assert manager._commands is mocked_create_registry.return_value
mocked_fetch_state.assert_called_once_with(["addr1", "addr2"])
mocked_create_registry.assert_called_once_with(
pooler.ensure_pool.return_value.execute.return_value)
assert pooler.ensure_pool.call_count == 3
assert pooler.ensure_pool.call_args_list[0] == mock.call(
Address("master1", 6666))
assert pooler.ensure_pool.call_args_list[1] == mock.call(
Address("master2", 7777))
assert pooler.ensure_pool.call_args_list[2] == mock.call(
Address("random_master", 5555))
await asyncio.sleep(0)
await manager.close()
def get_state():
return create_cluster_state(
SLOTS,
INFO,
Address("172.17.0.1", 6379),
)
async def test_execute__success_several_problem_retry(mocker):
cl = Cluster(["addr1"])
mocked_pooler = mocker.patch.object(cl, "_pooler", new=get_pooler_mock())
conn = mocked_pooler._conn
pool = mocked_pooler._pool
pool_execute_count = 0
errors = [
ConnectionRefusedError(),
ConnectionError(),
ConnectTimeoutError(("addr1", 6379)),
AskError("ASK 12539 1.2.3.4:9999"),
ClusterDownError("CLUSTERDOWN cluster is down"),
TryAgainError("TRYAGAIN try again later"),
MovedError("MOVED 12539 4.3.2.1:6666"),
ProtocolError(),
MovedError("MOVED 12539 4.3.2.1:6666"),
]
def conn_execute_se(*args, **kwargs):
nonlocal pool_execute_count
if args[0] == b"ASKING":
return conn.execute.return_value
pool_execute_count += 1
raise errors[pool_execute_count - 1]
def pool_execute_se(*args, **kwargs):
nonlocal pool_execute_count
pool_execute_count += 1
if pool_execute_count <= len(errors):
err = errors[pool_execute_count - 1]
raise err
return ("result", pool_execute_count)
pool.execute.side_effect = pool_execute_se
conn.execute.side_effect = conn_execute_se
mocked_manager = mocker.patch.object(cl,
"_manager",
new=get_manager_mock())
state = mocked_manager.get_state.return_value
mocked_execute_retry_slowdown = mocker.patch.object(
cl, "_execute_retry_slowdown", new=create_async_mock())
result = await cl.execute("get", "key")
assert mocked_manager.get_state.call_count == 10
assert mocked_pooler.ensure_pool.call_count == 10
assert mocked_manager.require_reload_state.call_count == 8
assert result == ("result", 10)
ensure_pool_calls = mocked_pooler.ensure_pool.call_args_list
# #1 attempt
assert ensure_pool_calls[0] == mock.call(
state.slot_master.return_value.addr)
# #2 attempt to random replica node
assert ensure_pool_calls[1] == mock.call(
state.random_slot_replica.return_value.addr)
# #3 attempt try to execute on random node in cluster
assert ensure_pool_calls[2] == mock.call(
state.random_node.return_value.addr)
# #4 attempt try random node in cluster
assert ensure_pool_calls[3] == mock.call(
state.random_node.return_value.addr)
# #5 attempt as ASKING to node from MovedError
assert ensure_pool_calls[4] == mock.call(Address("1.2.3.4", 9999))
# #6 attempt again random node because CLUSTERDOWN
assert ensure_pool_calls[5] == mock.call(
state.random_node.return_value.addr)
# #7 attempt again random node because TRYAGAIN
assert ensure_pool_calls[6] == mock.call(
state.random_node.return_value.addr)
# #8 attempt to node from MovedError
assert ensure_pool_calls[7] == mock.call(Address("4.3.2.1", 6666))
# #9 attempt again random node because ProtocolError
assert ensure_pool_calls[8] == mock.call(
state.random_node.return_value.addr)
# #10 attempt to node from MovedError, finally
assert ensure_pool_calls[9] == mock.call(Address("4.3.2.1", 6666))
# cl.determine_slot(b'key') == 12539
state.slot_master.assert_called_once_with(12539)
assert state.random_slot_replica.call_count == 1
assert state.random_slot_replica.call_args == mock.call(12539)
assert state.random_node.call_count == 5
assert mocked_execute_retry_slowdown.call_count == 9
asking_calls = [
c for c in conn.execute.call_args_list if c == mock.call(b"ASKING")
]
assert len(asking_calls) == 1
async def test_execute__error_after_max_attempts(mocker):
cl = Cluster(["addr1"])
mocked_pooler = mocker.patch.object(cl, "_pooler", new=get_pooler_mock())
pool = mocked_pooler._pool
pool_execute_count = 0
errors = [
MovedError("MOVED 12539 1.2.3.4:1000"),
MovedError("MOVED 12539 1.2.3.4:1001"),
MovedError("MOVED 12539 1.2.3.4:1002"),
MovedError("MOVED 12539 1.2.3.4:1003"),
MovedError("MOVED 12539 1.2.3.4:1004"),
MovedError("MOVED 12539 1.2.3.4:1005"),
MovedError("MOVED 12539 1.2.3.4:1006"),
MovedError("MOVED 12539 1.2.3.4:1007"),
MovedError("MOVED 12539 1.2.3.4:1008"),
MovedError("MOVED 12539 1.2.3.4:1009"),
MovedError("MOVED 12539 1.2.3.4:1010"),
MovedError("MOVED 12539 1.2.3.4:1011"),
]
def pool_execute_se(*args, **kwargs):
nonlocal pool_execute_count
pool_execute_count += 1
raise errors[pool_execute_count - 1]
pool.execute.side_effect = pool_execute_se
mocked_manager = mocker.patch.object(cl,
"_manager",
new=get_manager_mock())
state = mocked_manager.get_state.return_value
mocked_execute_retry_slowdown = mocker.patch.object(
cl, "_execute_retry_slowdown", new=create_async_mock())
with pytest.raises(MovedError) as ei:
await cl.execute("time")
assert (ei.value.info.host, ei.value.info.port) == ("1.2.3.4", 1009)
assert mocked_manager.get_state.call_count == 10
assert mocked_pooler.ensure_pool.call_count == 10
assert mocked_manager.require_reload_state.call_count == 10
ensure_pool_calls = mocked_pooler.ensure_pool.call_args_list
assert ensure_pool_calls[0] == mock.call(
state.random_master.return_value.addr)
assert ensure_pool_calls[1] == mock.call(Address("1.2.3.4", 1000))
assert ensure_pool_calls[2] == mock.call(Address("1.2.3.4", 1001))
assert ensure_pool_calls[3] == mock.call(Address("1.2.3.4", 1002))
assert ensure_pool_calls[4] == mock.call(Address("1.2.3.4", 1003))
assert ensure_pool_calls[5] == mock.call(Address("1.2.3.4", 1004))
assert ensure_pool_calls[6] == mock.call(Address("1.2.3.4", 1005))
assert ensure_pool_calls[7] == mock.call(Address("1.2.3.4", 1006))
assert ensure_pool_calls[8] == mock.call(Address("1.2.3.4", 1007))
assert ensure_pool_calls[9] == mock.call(Address("1.2.3.4", 1008))
assert state.random_node.call_count == 0
assert mocked_execute_retry_slowdown.call_count == 9
def create_cluster_state(
slots_resp: SlotsResponse,
cluster_info: Dict[str, str],
state_from: Address,
) -> ClusterState:
state_data = _ClusterStateData()
state_data.state = NodeClusterState(cluster_info[CLUSTER_INFO_STATE_KEY])
state_data.state_from = state_from
state_data.current_epoch = int(cluster_info[CLUSTER_INFO_CURRENT_EPOCH_KEY])
state_data.slots_assigned = int(cluster_info[CLUSTER_INFO_SLOTS_ASSIGNED])
nodes: Dict[Address, ClusterNode] = {}
master_replicas: Dict[Address, List[ClusterNode]] = {}
masters_set: Set[ClusterNode] = set()
replicas_set: Set[ClusterNode] = set()
for slot_spec in slots_resp:
master_node: Optional[ClusterNode] = None
slot_begin = int(slot_spec[0])
slot_end = int(slot_spec[1])
node_num = 0
for node_num, node_spec in enumerate(slot_spec[2:], 1):
node_addr = Address(node_spec[0], int(node_spec[1]))
if node_addr not in nodes:
node = ClusterNode(node_addr, node_spec[2])
nodes[node_addr] = node
else:
node = nodes[node_addr]
# first element is always master
if node_num == 1:
masters_set.add(node)
master_node = node
if master_node.addr not in master_replicas:
master_replicas[master_node.addr] = []
elif node_num > 1 and node not in replicas_set:
assert master_node is not None
replicas_set.add(node)
master_replicas[master_node.addr].append(node)
if node_num == 0:
logger.error(
"CLUSTER SLOTS returns slot range %r without master node",
(slot_begin, slot_end),
)
else:
assert master_node is not None
state_data.slots.append(
ClusterSlot(
begin=slot_begin,
end=slot_end,
master=master_node,
)
)
state_data.slots.sort(key=attrgetter("begin", "end"))
state_data.nodes = nodes
state_data.masters = list(masters_set)
state_data.replicas = master_replicas
state_data.addrs = list(nodes.keys())
return ClusterState(state_data)