fix: spop set index get error (#191)

* fix: spop set index get error * fix: spop set index get error * fix test case * fix test case * fix testcase --------- Co-authored-by: Tal Zaccai <talzacc@microsoft.com>
microsoft · Apr 4, 2024 · 04030dd · 04030dd
1 parent e72aab5
commit 04030dd
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 14 deletions.
diff --git a/libs/server/Objects/Set/SetObjectImpl.cs b/libs/server/Objects/Set/SetObjectImpl.cs
@@ -176,7 +176,7 @@ private void SetPop(byte* input, int length, ref SpanByteAndMemory output)
                     for (int i = 0; i < countParameter; i++)
                     {
                         // Generate a new index based on the elements left in the set
-                        var index = RandomNumberGenerator.GetInt32(0, set.Count - 1);
+                        var index = RandomNumberGenerator.GetInt32(0, set.Count);
                         var item = set.ElementAt(index);
                         set.Remove(item);
                         this.UpdateSize(item, false);

diff --git a/test/Garnet.test/RespSetTest.cs b/test/Garnet.test/RespSetTest.cs
@@ -368,26 +368,37 @@ public void CanDoSPOPCommandLC()
         [Test]
         public void CanDoSPOPWithCountCommandLC()
         {
-            var myset = new HashSet<string>();
-            myset.Add("one");
-            myset.Add("two");
-            myset.Add("three");
-            myset.Add("four");
-            myset.Add("five");
-
             CreateLongSet();
 
-            using var lightClientRequest = TestUtils.CreateRequest();
-            var response = lightClientRequest.SendCommand("SPOP myset 3", 4);
-            var strLen = Encoding.ASCII.GetString(response).Substring(1, 1);
-            Assert.AreEqual(3, Int32.Parse(strLen));
+            var lightClientRequest = TestUtils.CreateRequest();
+            var response = lightClientRequest.SendCommand("SPOP myset 3", 3);
+            var strResponse = Encoding.ASCII.GetString(response);
+            Assert.AreEqual('*', strResponse[0]);
+
+            var arrLenEndIdx = strResponse.IndexOf("\r\n", StringComparison.InvariantCultureIgnoreCase);
+            Assert.IsTrue(arrLenEndIdx > 1);
+
+            var strArrLen = Encoding.ASCII.GetString(response).Substring(1, arrLenEndIdx - 1);
+            Assert.IsTrue(int.TryParse(strArrLen, out var arrLen));
+            Assert.AreEqual(3, arrLen);
 
             var secondResponse = lightClientRequest.SendCommands("SCARD myset", "PING", 1, 1);
             var expectedResponse = ":2\r\n+PONG\r\n";
-            var strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
+            strResponse = Encoding.ASCII.GetString(secondResponse).Substring(0, expectedResponse.Length);
             Assert.AreEqual(expectedResponse, strResponse);
-        }
 
+            // Test for popping set until empty
+            response = lightClientRequest.SendCommand("SPOP myset 2", 2);
+            strResponse = Encoding.ASCII.GetString(response);
+            Assert.AreEqual('*', strResponse[0]);
+
+            arrLenEndIdx = strResponse.IndexOf("\r\n", StringComparison.InvariantCultureIgnoreCase);
+            Assert.IsTrue(arrLenEndIdx > 1);
+
+            strArrLen = Encoding.ASCII.GetString(response).Substring(1, arrLenEndIdx - 1);
+            Assert.IsTrue(int.TryParse(strArrLen, out arrLen));
+            Assert.AreEqual(2, arrLen);
+        }
 
         [Test]
         public void MultiWithNonExistingSet()