ClickHouse/tests/integration/test_cluster_discovery/test.py
2022-03-22 17:39:58 +01:00

94 lines
2.5 KiB
Python

import pytest
import functools
import time
from helpers.cluster import ClickHouseCluster
cluster = ClickHouseCluster(__file__)
shard_configs = {i: f"config/config_shard{i}.xml" for i in [1, 3]}
nodes = [
cluster.add_instance(
f"node{i}",
main_configs=[shard_configs.get(i, "config/config.xml")],
stay_alive=True,
with_zookeeper=True,
)
for i in range(5)
]
@pytest.fixture(scope="module")
def start_cluster():
try:
cluster.start()
yield cluster
finally:
cluster.shutdown()
def check_on_cluster(
nodes, expected, *, what, cluster_name="test_auto_cluster", msg=None, retries=5
):
"""
Select data from `system.clusters` on specified nodes and check the result
"""
assert 1 <= retries <= 6
for retry in range(1, retries + 1):
nodes_res = {
node.name: int(
node.query(
f"SELECT {what} FROM system.clusters WHERE cluster = '{cluster_name}'"
)
)
for node in nodes
}
if all(actual == expected for actual in nodes_res.values()):
break
if retry != retries:
time.sleep(2**retry)
else:
msg = msg or f"Wrong '{what}' result"
raise Exception(
f"{msg}: {nodes_res}, expected: {expected} (after {retries} retries)"
)
def test_cluster_discovery_startup_and_stop(start_cluster):
"""
Start cluster, check nodes count in system.clusters,
then stop/start some nodes and check that it (dis)appeared in cluster.
"""
check_nodes_count = functools.partial(
check_on_cluster, what="count()", msg="Wrong nodes count in cluster"
)
check_shard_num = functools.partial(
check_on_cluster,
what="count(DISTINCT shard_num)",
msg="Wrong shard_num count in cluster",
)
total_shards = len(shard_configs) + 1
check_nodes_count([nodes[0], nodes[2]], len(nodes))
check_shard_num([nodes[0], nodes[2]], total_shards)
nodes[1].stop_clickhouse(kill=True)
check_nodes_count([nodes[0], nodes[2]], len(nodes) - 1)
check_shard_num([nodes[0], nodes[2]], total_shards - 1)
nodes[3].stop_clickhouse()
check_nodes_count([nodes[0], nodes[2]], len(nodes) - 2)
nodes[1].start_clickhouse()
check_nodes_count([nodes[0], nodes[2]], len(nodes) - 1)
nodes[3].start_clickhouse()
check_nodes_count([nodes[0], nodes[2]], len(nodes))
check_nodes_count([nodes[1], nodes[2]], 2, cluster_name="two_shards", retries=1)