From b924b574ea0f895671d535b96b387eb8f591f151 Mon Sep 17 00:00:00 2001 From: Tracy Date: Wed, 10 Jun 2026 22:31:30 +0000 Subject: [PATCH 1/5] Add Cluster Mode Enabled (CME) testing support Signed-off-by: Tracy --- src/valkey_test_case.py | 268 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 268 insertions(+) diff --git a/src/valkey_test_case.py b/src/valkey_test_case.py index 90a3470..d85c76a 100644 --- a/src/valkey_test_case.py +++ b/src/valkey_test_case.py @@ -728,3 +728,271 @@ def waitForReplicaOffsetToSyncUp(self, primary, replica): pinfo.get_primary_repl_offset(), timeout=TEST_MAX_WAIT_TIME_SECONDS, ) + + +from valkey.cluster import VALKEY_CLUSTER_HASH_SLOTS + + +class ClusterInfo: + """Contains information about a point in time of a Valkey cluster.""" + + def __init__(self, info): + self.info = info + + def get_cluster_my_epoch(self): + return self.info["cluster_my_epoch"] + + def get_cluster_epoch(self): + return self.info["cluster_current_epoch"] + + def is_cluster_ok(self): + """Return True if the cluster state is OK.""" + return self.info["cluster_state"] == "ok" + + def is_cluster_down(self): + """Return True if the cluster state is fail.""" + return self.info["cluster_state"] == "fail" + + 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 start(self, connect_client=True): + super(ClusterNodeHandle, self).start(connect_client=connect_client) + if connect_client: + self._set_node_id() + + 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_know_node(self, nodeid): + def knows(): + nodesInfo = self.client.cluster("NODES") + for key in nodesInfo: + if re.match(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): + for node in self.nodes: + node.start() + + 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): + 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_know_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. + try: + self.nodes[i].client.readonly() + except Exception: + logging.warning( + "READONLY failed on replica port {}".format(self.nodes[i].port) + ) + + 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.""" + from valkey.cluster import ValkeyCluster + + 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) From c66deb546001108acaa620511031a6a9afcf8b76 Mon Sep 17 00:00:00 2001 From: Tracy Date: Thu, 11 Jun 2026 18:13:14 +0000 Subject: [PATCH 2/5] ADD CME integration tests and remove some unused functions Signed-off-by: Tracy --- src/valkey_test_case.py | 16 +---- tests/test_cluster_mode.py | 123 +++++++++++++++++++++++++++++++++++++ 2 files changed, 125 insertions(+), 14 deletions(-) create mode 100644 tests/test_cluster_mode.py diff --git a/src/valkey_test_case.py b/src/valkey_test_case.py index d85c76a..505d562 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 from enum import Enum MAX_PING_TRIES = 60 @@ -730,29 +730,16 @@ def waitForReplicaOffsetToSyncUp(self, primary, replica): ) -from valkey.cluster import VALKEY_CLUSTER_HASH_SLOTS - - class ClusterInfo: """Contains information about a point in time of a Valkey cluster.""" def __init__(self, info): self.info = info - def get_cluster_my_epoch(self): - return self.info["cluster_my_epoch"] - - def get_cluster_epoch(self): - return self.info["cluster_current_epoch"] - def is_cluster_ok(self): """Return True if the cluster state is OK.""" return self.info["cluster_state"] == "ok" - def is_cluster_down(self): - """Return True if the cluster state is fail.""" - return self.info["cluster_state"] == "fail" - def cluster_known_nodes(self): """Return the number of nodes known to this node.""" return int(self.info["cluster_known_nodes"]) @@ -907,6 +894,7 @@ def create_cluster(self, num_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)) 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 From 2e09c8a54ec2f5541380e2759041b01dbef39e9a Mon Sep 17 00:00:00 2001 From: Tracy Date: Wed, 1 Jul 2026 18:15:51 +0000 Subject: [PATCH 3/5] Fix nits from CME review Signed-off-by: Tracy --- src/valkey_test_case.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/src/valkey_test_case.py b/src/valkey_test_case.py index 505d562..3704fae 100644 --- a/src/valkey_test_case.py +++ b/src/valkey_test_case.py @@ -783,11 +783,6 @@ def _set_node_id(self): ) return - def start(self, connect_client=True): - super(ClusterNodeHandle, self).start(connect_client=connect_client) - if connect_client: - self._set_node_id() - def connect(self): client = super(ClusterNodeHandle, self).connect() self._set_node_id() @@ -821,11 +816,11 @@ def wait_for_cluster_known_nodes(self, count): timeout=TEST_MAX_WAIT_TIME_SECONDS, ) - def wait_for_cluster_know_node(self, nodeid): + def wait_for_cluster_known_node(self, nodeid): def knows(): nodesInfo = self.client.cluster("NODES") for key in nodesInfo: - if re.match(nodeid, nodesInfo[key]["node_id"]): + if nodeid == nodesInfo[key]["node_id"]: return True return False @@ -875,9 +870,9 @@ def create_nodes(self, num_nodes): self.create_node() return self.nodes - def start_all_nodes(self): + def start_all_nodes(self, wait_for_ping=True, connect_client=True): for node in self.nodes: - node.start() + 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.""" @@ -908,7 +903,7 @@ def setup_replicas(self, num_shards, num_replicas_per_shard): 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_know_node(primary.nodeid) + 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. From 90c3ae23948710ff34a5041f89f4555c6765622f Mon Sep 17 00:00:00 2001 From: Tracy Date: Wed, 1 Jul 2026 18:34:25 +0000 Subject: [PATCH 4/5] Fix nits Signed-off-by: Tracy --- src/valkey_test_case.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/src/valkey_test_case.py b/src/valkey_test_case.py index 3704fae..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 +from valkey.cluster import VALKEY_CLUSTER_HASH_SLOTS, ValkeyCluster from enum import Enum MAX_PING_TRIES = 60 @@ -916,12 +916,7 @@ def setup_replicas(self, num_shards, num_replicas_per_shard): for i in range(num_shards, total): self.waitForReplicaToSyncUp(self.nodes[i]) # Allow read-only queries to be served by the replica. - try: - self.nodes[i].client.readonly() - except Exception: - logging.warning( - "READONLY failed on replica port {}".format(self.nodes[i].port) - ) + 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. @@ -950,8 +945,6 @@ def setup_cluster(self, num_shards, num_replicas_per_shard): def get_cluster_client(self): """Return a cluster-aware client that follows MOVED/ASK redirections.""" - from valkey.cluster import ValkeyCluster - 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 From 34f391ffcf0d30a2e5758b06b4f77acc11d72de7 Mon Sep 17 00:00:00 2001 From: Tracy Date: Thu, 2 Jul 2026 00:02:32 +0000 Subject: [PATCH 5/5] Update readme Signed-off-by: Tracy --- README.md | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) 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`.