From 6b274dfdcc49573a4f5021b710a8c2bdd75ab7f0 Mon Sep 17 00:00:00 2001
From: ananyabalagere <abalagere@gmail.com>
Date: Fri, 14 Mar 2025 00:24:02 -0700
Subject: [PATCH] shuffle two

---
 tests/shuffle/basic_shuffle.py | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/tests/shuffle/basic_shuffle.py b/tests/shuffle/basic_shuffle.py
index 7e8dcb5..836a2a7 100644
--- a/tests/shuffle/basic_shuffle.py
+++ b/tests/shuffle/basic_shuffle.py
@@ -122,13 +122,32 @@ def basic_shuffle_2(conductor: ClusterConductor, dir, log: Logger):
             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 % 4
+            node_to_put = node_to_put % 3
 
         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}"
         original_get_all = r.json()["items"]
 
         conductor.add_shard("shard2", conductor.get_nodes([2]))
+        fx.broadcast_view(conductor.get_shard_view())
+
+        r = c.get_all(2, timeout=10) # should get all of shard 2's keys
+        assert r.ok, f"expected ok for get, got {r.status_code}"
+        get_all_1 = r.json()["items"]
+        keys1 = get_all_1.keys()
+
+        r = c.get_all(1, timeout=10) # should get all of shard 1's keys
+        assert r.ok, f"expected ok for get, got {r.status_code}"
+        get_all_2 = r.json()["items"]
+        keys2 = get_all_2.keys()
+
+        for key in keys1:
+            assert not (key in keys2)
+        for key in keys2:
+            assert not (key in keys1)
+        
+        assert original_get_all.keys() == keys1 + keys2
+
         return True, "ok"
 
-SHUFFLE_TESTS = [TestCase("shuffle_basic", basic_shuffle)]
\ No newline at end of file
+SHUFFLE_TESTS = [TestCase("basic_shuffle", basic_shuffle), TestCase("basic_shuffle_2", basic_shuffle_2)]
\ No newline at end of file
-- 
GitLab