diff --git a/README.md b/README.md index acd6a56..3fa2218 100644 --- a/README.md +++ b/README.md @@ -67,4 +67,26 @@ class TestExamplePerClassSetup(ExampleTestCaseBase): client.execute_command("SET K V") ``` -For more examples, refer to the `tests` directory of this package. +**Testing Cluster Mode Enabled (CME)** + +To test against a Valkey cluster, inherit `ClusterTestCase` instead of `ValkeyTestCase`. Call `setup_cluster(num_shards, num_replicas_per_shard)` to bootstrap a cluster: it starts all nodes, assigns slots, attaches replicas, waits until the cluster state is `ok`, and returns a client that follows `MOVED`/`ASK` redirections. The cluster is torn down automatically after each test. + +``` +from valkey_test_case import ClusterTestCase, ClusterInfo + +class TestExampleCluster(ClusterTestCase): + def test_cluster_read_write(self): + self.server_path = "/path_to_your_valkey_server_binary" + + # 3 primaries, each with 1 replica (6 nodes total) + client = self.setup_cluster(num_shards=3, num_replicas_per_shard=1) + + client.set("key", "value") + assert client.get("key") == b"value" + + # self.nodes holds every ClusterNodeHandle for direct inspection + for node in self.nodes: + assert ClusterInfo(node.client.cluster("INFO")).is_cluster_ok() +``` + +To apply startup arguments (modules, configs) to every node, set `self.args` inside the test before calling `setup_cluster`. diff --git a/src/valkey_test_case.py b/src/valkey_test_case.py index 90a3470..8e01064 100644 --- a/src/valkey_test_case.py +++ b/src/valkey_test_case.py @@ -8,7 +8,7 @@ from functools import wraps from valkey import * from util.waiters import * - +from valkey.cluster import VALKEY_CLUSTER_HASH_SLOTS, ValkeyCluster from enum import Enum MAX_PING_TRIES = 60 @@ -728,3 +728,247 @@ def waitForReplicaOffsetToSyncUp(self, primary, replica): pinfo.get_primary_repl_offset(), timeout=TEST_MAX_WAIT_TIME_SECONDS, ) + + +class ClusterInfo: + """Contains information about a point in time of a Valkey cluster.""" + + def __init__(self, info): + self.info = info + + def is_cluster_ok(self): + """Return True if the cluster state is OK.""" + return self.info["cluster_state"] == "ok" + + def cluster_known_nodes(self): + """Return the number of nodes known to this node.""" + return int(self.info["cluster_known_nodes"]) + + def cluster_slots_assigned(self): + """Return the number of hash slots currently assigned.""" + return int(self.info["cluster_slots_assigned"]) + + +class ClusterNodeHandle(ValkeyServerHandle): + """Handle to a valkey server process running in cluster mode enabled (CME).""" + + def __init__( + self, + bind_ip, + port, + port_tracker, + testdir, + server_path="valkey-server", + ): + super(ClusterNodeHandle, self).__init__( + bind_ip, port, port_tracker, server_path=server_path, cwd=testdir + ) + # Start the node in cluster mode. The cluster-config-file (nodes.conf) + # is per-node and cleaned up by the base class teardown. + self.args["cluster-enabled"] = "yes" + self.args["cluster-config-file"] = "nodes_{}_{}.conf".format(bind_ip, port) + self.args["cluster-node-timeout"] = "2000" + self.masterid = None + self.nodeid = None + + def _set_node_id(self): + # CLUSTER NODES should return only one node - myself after startup + # Read the node id + nodes = self.client.cluster("NODES") + for key in nodes: + if re.match("myself", nodes[key]["flags"]): + self.nodeid = nodes[key]["node_id"] + logging.info( + "Cluster node {} has node id {}".format(self.port, self.nodeid) + ) + return + + def connect(self): + client = super(ClusterNodeHandle, self).connect() + self._set_node_id() + return client + + def meet(self, ip, port): + return self.client.execute_command("CLUSTER", "MEET", ip, port) + + def replicate(self, node_id): + self.masterid = node_id + return self.client.execute_command("CLUSTER", "REPLICATE", node_id) + + def assign_slots(self, *args): + """ + Assign multiple ranges of slots to this node. + Accepts multiple pairs that are interpretted as [low,high) + assign_slots(0,10) -> [0..9] + assign_slots(0,10,15,20) -> [0..9] and [15..19] + """ + assert len(args) % 2 == 0 + command = ["CLUSTER", "ADDSLOTSRANGE"] + for t in range(0, len(args), 2): + command.extend([args[t], args[t + 1] - 1]) + return self.client.execute_command(*command) + + def wait_for_cluster_known_nodes(self, count): + """Wait until we are connected to exactly count nodes.""" + wait_for_equal( + lambda: ClusterInfo(self.client.cluster("INFO")).cluster_known_nodes(), + count, + timeout=TEST_MAX_WAIT_TIME_SECONDS, + ) + + def wait_for_cluster_known_node(self, nodeid): + def knows(): + nodesInfo = self.client.cluster("NODES") + for key in nodesInfo: + if nodeid == nodesInfo[key]["node_id"]: + return True + return False + + wait_for_true(knows, timeout=TEST_MAX_WAIT_TIME_SECONDS) + + def wait_for_cluster_ok(self): + wait_for_true( + lambda: ClusterInfo(self.client.cluster("INFO")).is_cluster_ok(), + timeout=TEST_MAX_WAIT_TIME_SECONDS, + ) + + +class ClusterTestCase(ValkeyTestCase): + """Base class for Cluster Mode Enabled (CME) tests.""" + + @pytest.fixture(autouse=True) + def cluster_setup(self): + # Per-test cluster state. Initialized here rather than as class-level + # attributes so each test starts with its own node list. + self.nodes = [] + self.cluster_client = None + yield + self.teardown() + + def create_node(self, bind_ip=None, port=None): + """Create a single cluster-mode node and register it for teardown.""" + if not bind_ip: + bind_ip = self.get_bind_ip() + if not port: + port = self.get_bind_port() + + node = ClusterNodeHandle( + bind_ip=bind_ip, + port=port, + port_tracker=self.port_tracker, + testdir=self.testdir, + server_path=self.server_path, + ) + node.args.update(self.args) + self.nodes.append(node) + # Registered in server_list so ValkeyTestCase.teardown reclaims it. + self.server_list.append(node) + return node + + def create_nodes(self, num_nodes): + for _ in range(num_nodes): + self.create_node() + return self.nodes + + def start_all_nodes(self, wait_for_ping=True, connect_client=True): + for node in self.nodes: + node.start(wait_for_ping=wait_for_ping, connect_client=connect_client) + + def create_cluster(self, num_nodes): + """Start `num_nodes` nodes and gossip them into a single cluster.""" + self.create_nodes(num_nodes) + self.start_all_nodes() + + # Introduce every other node to the first node; gossip propagates the + # full topology from there. + for i in range(1, num_nodes): + self.nodes[0].meet(self.nodes[i].bind_ip, self.nodes[i].port) + + # Wait until every node has discovered the whole cluster. + for node in self.nodes: + node.wait_for_cluster_known_nodes(num_nodes) + + def assign_slots_to_nodes(self, num_nodes): + """Equally distribute sequential slots to each node.""" + slot_slice = VALKEY_CLUSTER_HASH_SLOTS / num_nodes + for i in range(num_nodes): + slot_min = int(round(slot_slice * i)) + slot_max = int(round(slot_slice * (i + 1))) + self.nodes[i].assign_slots(slot_min, slot_max) + + def setup_replicas(self, num_shards, num_replicas_per_shard): + """Attach the remaining nodes as replicas, round-robin across shards.""" + total = num_shards * (1 + num_replicas_per_shard) + for i in range(num_shards, total): + shard_idx = i % num_shards + primary = self.nodes[shard_idx] + # Make sure the replica knows the primary before replicating. + self.nodes[i].wait_for_cluster_known_node(primary.nodeid) + self.nodes[i].replicate(primary.nodeid) + + # Wait for each shard's replicas to come online and sync up. + for i in range(num_shards): + wait_for_equal( + lambda primary=self.nodes[i]: primary.num_replicas_online(), + num_replicas_per_shard, + timeout=MAX_REPLICA_WAIT_TIME, + ) + for i in range(num_shards, total): + self.waitForReplicaToSyncUp(self.nodes[i]) + # Allow read-only queries to be served by the replica. + self.nodes[i].client.readonly() + + def setup_cluster(self, num_shards, num_replicas_per_shard): + """Create and fully bootstrap a cluster, returning a cluster client. + + When this returns the cluster is in the 'ok' state and ready to serve. + """ + total_nodes = num_shards * (1 + num_replicas_per_shard) + self.create_cluster(total_nodes) + + self.assign_slots_to_nodes(num_shards) + + # Bump each primary's config epoch so replicas don't overtake it via + # epoch collision resolution. + for i in range(num_shards): + self.nodes[i].client.execute_command("CLUSTER", "BUMPEPOCH") + + if num_replicas_per_shard > 0: + self.setup_replicas(num_shards, num_replicas_per_shard) + + # Wait for every node to agree the cluster is healthy. + for node in self.nodes: + node.wait_for_cluster_ok() + + self.cluster_client = self.get_cluster_client() + return self.cluster_client + + def get_cluster_client(self): + """Return a cluster-aware client that follows MOVED/ASK redirections.""" + primary = self.nodes[0] + # A cluster client discovers the topology by connecting to the host each + # node advertises in CLUSTER SLOTS, not the bind address. When a node + # binds the wildcard 0.0.0.0 it advertises a concrete, connectable host + # (typically loopback) instead, so read that advertised host back and + # seed the client with it rather than the bind address. + host = primary.bind_ip + for slot_range in primary.client.cluster("SLOTS"): + # slot_range = [start, end, [host, port, node_id, ...], ...] + owner_host, _, owner_id = ( + slot_range[2][0], + slot_range[2][1], + slot_range[2][2], + ) + if owner_id.decode() == primary.nodeid: + host = owner_host.decode() + break + return ValkeyCluster(host=host, port=primary.port) + + def teardown(self): + if self.cluster_client is not None: + try: + self.cluster_client.close() + except Exception: + pass + self.cluster_client = None + ValkeyTestCase.teardown(self) diff --git a/tests/test_cluster_mode.py b/tests/test_cluster_mode.py new file mode 100644 index 0000000..e51fe08 --- /dev/null +++ b/tests/test_cluster_mode.py @@ -0,0 +1,123 @@ +from conftest import resource_port_tracker +from valkey_test_case import ClusterTestCase, ClusterInfo +import pytest +import os + + +SERVER_PATH = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + ".build", + "binaries", + os.environ["SERVER_VERSION"], + "valkey-server", +) + + +class TestClusterBasic(ClusterTestCase): + """Verify that ClusterTestCase can bootstrap a 3-primary cluster.""" + + server_path = SERVER_PATH + + @pytest.mark.order(1) + def test_setup_cluster_and_use(self): + client = self.setup_cluster(num_shards=3, num_replicas_per_shard=0) + + # Verify cluster_state is ok on all nodes + for node in self.nodes: + assert ClusterInfo( + node.client.cluster("INFO") + ).is_cluster_ok(), f"Node {node.port} not in ok state" + + # Verify all 16384 slots are assigned + info = ClusterInfo(self.nodes[0].client.cluster("INFO")) + assert info.cluster_slots_assigned() == 16384 + assert info.cluster_known_nodes() == 3 + + # Write and read keys — ValkeyCluster follows MOVED redirections + for i in range(100): + client.set(f"key:{i}", f"value:{i}") + + for i in range(100): + val = client.get(f"key:{i}") + assert val == b"value:" + str(i).encode() + + @pytest.mark.order(2) + @pytest.mark.parametrize("num_shards", [2, 3, 5]) + def test_slot_distribution(self, num_shards): + self.setup_cluster(num_shards=num_shards, num_replicas_per_shard=0) + + # Collect the set of slots each primary owns from CLUSTER SLOTS. + # Reply shape: [[start, end, [host, port, node_id, ...], ...], ...] + owned = {} + for slot_range in self.nodes[0].client.cluster("SLOTS"): + start, end = slot_range[0], slot_range[1] + owner_id = slot_range[2][2].decode() + owned.setdefault(owner_id, set()).update(range(start, end + 1)) + + # Every primary owns a share of the slots. + assert len(owned) == num_shards + + all_slots = set() + for slots in owned.values(): + # No slot is claimed by more than one primary. + assert all_slots.isdisjoint(slots) + all_slots |= slots + + # Every one of the 16384 slots is covered exactly once. + assert all_slots == set(range(16384)) + + +class TestClusterWithReplicas(ClusterTestCase): + """Verify multi-shard clusters with replicas.""" + + server_path = SERVER_PATH + + @pytest.mark.order(1) + def test_setup_cluster_with_replicas(self): + client = self.setup_cluster(num_shards=3, num_replicas_per_shard=1) + + # 6 nodes total + assert len(self.nodes) == 6 + + # Cluster is healthy + for node in self.nodes: + assert ClusterInfo( + node.client.cluster("INFO") + ).is_cluster_ok(), f"Node {node.port} not in ok state" + + # All slots assigned + info = ClusterInfo(self.nodes[0].client.cluster("INFO")) + assert info.cluster_slots_assigned() == 16384 + assert info.cluster_known_nodes() == 6 + + # Replicas have master IDs set + for i in range(3, 6): + assert self.nodes[i].masterid is not None + + # Write/read through cluster client + for i in range(50): + client.set(f"rkey:{i}", f"rval:{i}") + for i in range(50): + assert client.get(f"rkey:{i}") == f"rval:{i}".encode() + + @pytest.mark.order(2) + def test_setup_cluster_multiple_replicas_per_shard(self): + # 2 shards x (1 primary + 2 replicas) = 6 nodes + self.setup_cluster(num_shards=2, num_replicas_per_shard=2) + + assert len(self.nodes) == 6 + + # Cluster is healthy and every node sees all 6 + for node in self.nodes: + info = ClusterInfo(node.client.cluster("INFO")) + assert info.is_cluster_ok(), f"Node {node.port} not in ok state" + assert info.cluster_known_nodes() == 6 + + # Each primary has exactly 2 replicas online. + for i in range(2): + assert self.nodes[i].num_replicas_online() == 2 + + # Every replica is attached to one of the two primaries. + primary_ids = {self.nodes[0].nodeid, self.nodes[1].nodeid} + for i in range(2, 6): + assert self.nodes[i].masterid in primary_ids