diff --git a/tests/shuffle/basic_shuffle.py b/tests/shuffle/basic_shuffle.py index 3272be2739ab33eeb1780c7246b20bc5b25f926b..898061c4ecf6a810f48b2ac6daf06a9926dd2530 100644 --- a/tests/shuffle/basic_shuffle.py +++ b/tests/shuffle/basic_shuffle.py @@ -4,6 +4,102 @@ from ...utils.util import Logger from ..helper import KVSMultiClient, KVSTestFixture from ...utils.kvs_api import DEFAULT_TIMEOUT +def basic_shuffle_add_remove(conductor: ClusterConductor, dir, log: Logger): + with KVSTestFixture(conductor, dir, log, node_count=3) as fx: + c = KVSMultiClient(fx.clients, "client", log) + conductor.add_shard("shard1", conductor.get_nodes([0])) + conductor.add_shard("shard2", conductor.get_nodes([1])) + + fx.broadcast_view(conductor.get_shard_view()) + + node_to_put = 0 + base_key = "key" + # Put 15 keys + for i in range(15): + log(f"Putting key {i}\n") + r = c.put(node_to_put, f"{base_key}{i}", f"{i}", timeout=10) + assert r.ok, f"expected ok for new key, got {r.status_code}" + node_to_put += 1 + node_to_put = node_to_put % 2 + + # Get all keys + r = c.get_all(0, timeout=10) + assert r.ok, f"expected ok for get, got {r.status_code}" + res = r.json()["items"] + shard1_keys = res + + r = c.get_all(1, timeout=10) + assert r.ok, f"expected ok for get, got {r.status_code}" + res = r.json()["items"] + shard2_keys = res + + log(f"Shard 1 keys: {shard1_keys}\n") + log(f"Shard 2 keys: {shard2_keys}\n") + + # Total number of keys should matched number of keys put + assert len(shard1_keys) + len(shard2_keys) == 15, f"expected 15 keys, got {len(shard1_keys) + len(shard2_keys)}" + + # Add a 3rd shard, causing a shuffle. There should still be 15 keys at the end. + log("Adding 3rd shard\n") + conductor.add_shard("shard3", conductor.get_nodes([2])) + fx.broadcast_view(conductor.get_shard_view()) + + # Get the keys on shard 1 + r = c.get_all(0, timeout=10) + assert r.ok, f"expected ok for get, got {r.status_code}" + res = r.json()["items"] + shard1_keys = res + + log(f"Shard 1 keys: {shard1_keys}\n") + + # get the keys on shard 2 + r = c.get_all(1, timeout=10) + assert r.ok, f"expected ok for get, got {r.status_code}" + res = r.json()["items"] + shard2_keys = res + + log(f"Shard 2 keys: {shard2_keys}\n") + + # get the keys on shard 3 + r = c.get_all(2, timeout=10) + assert r.ok, f"expected ok for get, got {r.status_code}" + res = r.json()["items"] + shard3_keys = res + + log(f"Shard 3 keys: {shard3_keys}\n") + + assert len(shard1_keys) + len(shard2_keys) + len(shard3_keys) == 15, f"expected 15 keys, got {len(shard1_keys) + len(shard2_keys) + len(shard3_keys)}" + + # Remove shard 3, causing a shuffle. Move Node 2 to shard 1 so the keys should still exist, and be shuffled + conductor.remove_shard("shard3") + fx.broadcast_view(conductor.get_shard_view()) + + r = c.get_all(0, timeout=10) # should get all of shard 1's keys + assert r.ok, f"expected ok for get, got {r.status_code}" + res = r.json()["items"] + shard1_keys = res + + r = c.get_all(1, timeout=10) # should get all of shard 2's keys + assert r.ok, f"expected ok for get, got {r.status_code}" + res = r.json()["items"] + shard2_keys = res + + assert len(shard1_keys) + len(shard2_keys) == 15, f"expected 15 keys, got {len(shard1_keys) + len(shard2_keys)}" + + # Remove shard 2. This loses keys. + conductor.remove_shard("shard2") + fx.broadcast_view(conductor.get_shard_view()) + + r = c.get_all(0, timeout=10) + assert r.ok, f"expected ok for get, got {r.status_code}" + res = r.json()["items"] + shard1_keys_after_delete = res + + assert len(shard1_keys_after_delete) == 15, f"expected 15 keys, got {len(shard1_keys_after_delete)}" + + + return True, "ok" + def basic_shuffle(conductor: ClusterConductor, dir, log: Logger): with KVSTestFixture(conductor, dir, log, node_count=3) as fx: c = KVSMultiClient(fx.clients, "client", log) @@ -89,6 +185,7 @@ def basic_shuffle(conductor: ClusterConductor, dir, log: Logger): # Remove shard 2. This loses keys. conductor.remove_shard("shard2") + conductor.remove_node_from_shard("shard1", conductor.get_node(2)) fx.broadcast_view(conductor.get_shard_view()) r = c.get_all(0, timeout=10) @@ -96,7 +193,7 @@ def basic_shuffle(conductor: ClusterConductor, dir, log: Logger): res = r.json()["items"] shard1_keys_after_delete = res - assert len(shard1_keys) == len(shard1_keys_after_delete), f"expected {len(shard1_keys)} keys, got {len(shard1_keys_after_delete)}" + assert len(shard1_keys_after_delete) == 15, f"expected 15 keys, got {len(shard1_keys_after_delete)}" return True, "ok" @@ -151,4 +248,4 @@ def basic_shuffle_2(conductor: ClusterConductor, dir, log: Logger): return True, "ok" -SHUFFLE_TESTS = [TestCase("basic_shuffle", basic_shuffle), TestCase("basic_shuffle_2", basic_shuffle_2)] \ No newline at end of file +SHUFFLE_TESTS = [TestCase("basic_shuffle_add_remove", basic_shuffle_add_remove), TestCase("basic_shuffle", basic_shuffle), TestCase("basic_shuffle_2", basic_shuffle_2)]