Added Equals receiver funcion on SortedSet to compare deep equality of two sorted sets. Created unit test for ZMPOP command handler

src/modules/sorted_set/commands.go

@@ -628,8 +628,9 @@ func handleZINTERSTORE(ctx context.Context, cmd []string, server utils.Server, c
 }
 
 func handleZMPOP(ctx context.Context, cmd []string, server utils.Server, conn *net.Conn) ([]byte, error) {
-	if len(cmd) < 2 {
+	keys, err := zmpopKeyFunc(cmd)
-		return nil, errors.New(utils.WRONG_ARGS_RESPONSE)
+	if err != nil {
+		return nil, err
 	}
 
 	count := 1
@@ -672,48 +673,36 @@ func handleZMPOP(ctx context.Context, cmd []string, server utils.Server, conn *n
 		}
 	}
 
-	var keys []string
+	for i := 0; i < len(keys); i++ {
-	if modifierIdx == -1 {
+		if server.KeyExists(keys[i]) {
-		keys = cmd[1:]
+			if _, err = server.KeyLock(ctx, keys[i]); err != nil {
-	} else {
-		keys = cmd[1:modifierIdx]
-	}
-
-	for _, key := range keys {
-		if server.KeyExists(key) {
-			_, err := server.KeyLock(ctx, key)
-			if err != nil {
 				continue
 			}
-			v, ok := server.GetValue(key).(*SortedSet)
+			v, ok := server.GetValue(keys[i]).(*SortedSet)
 			if !ok || v.Cardinality() == 0 {
-				server.KeyUnlock(key)
+				server.KeyUnlock(keys[i])
 				continue
 			}
 			popped, err := v.Pop(count, policy)
 			if err != nil {
-				server.KeyUnlock(key)
+				server.KeyUnlock(keys[i])
 				return nil, err
 			}
-			server.KeyUnlock(key)
+			server.KeyUnlock(keys[i])
-			if popped.Cardinality() == 0 {
-				return []byte("+(nil)\r\n\r\n"), nil
-			}
 
 			res := fmt.Sprintf("*%d", popped.Cardinality())
-			for i, m := range popped.GetAll() {
+
-				s := fmt.Sprintf("%s %f", m.value, m.score)
+			for _, m := range popped.GetAll() {
-				res += fmt.Sprintf("\r\n$%d\r\n%s", len(s), s)
+				res += fmt.Sprintf("\r\n*2\r\n$%d\r\n%s\r\n+%s", len(m.value), m.value, strconv.FormatFloat(float64(m.score), 'f', -1, 64))
-				if i == popped.Cardinality()-1 {
+			}
+
 			res += "\r\n\r\n"
-				}
-			}
 
 			return []byte(res), nil
 		}
 	}
 
-	return []byte("+(nil)\r\n\r\n"), nil
+	return []byte("*0\r\n\r\n"), nil
 }
 
 func handleZPOP(ctx context.Context, cmd []string, server utils.Server, conn *net.Conn) ([]byte, error) {

src/modules/sorted_set/commands_test.go

@@ -1162,6 +1162,244 @@ func Test_HandleZINCRBY(t *testing.T) {
 	}
 }
 
+func Test_HandleZMPOP(t *testing.T) {
+	mockServer := server.NewServer(server.Opts{})
+
+	tests := []struct {
+		preset           bool
+		presetValues     map[string]interface{}
+		command          []string
+		expectedValues   map[string]*SortedSet
+		expectedResponse [][]string
+		expectedError    error
+	}{
+		{ // 1. Successfully pop one min element by default
+			preset: true,
+			presetValues: map[string]interface{}{
+				"key1": NewSortedSet([]MemberParam{
+					{value: "one", score: 1}, {value: "two", score: 2},
+					{value: "three", score: 3}, {value: "four", score: 4},
+					{value: "five", score: 5},
+				}),
+			},
+			command: []string{"ZMPOP", "key1"},
+			expectedValues: map[string]*SortedSet{
+				"key1": NewSortedSet([]MemberParam{
+					{value: "two", score: 2},
+					{value: "three", score: 3}, {value: "four", score: 4},
+					{value: "five", score: 5},
+				}),
+			},
+			expectedResponse: [][]string{
+				{"one", "1"},
+			},
+			expectedError: nil,
+		},
+		{ // 2. Successfully pop one min element by specifying MIN
+			preset: true,
+			presetValues: map[string]interface{}{
+				"key2": NewSortedSet([]MemberParam{
+					{value: "one", score: 1}, {value: "two", score: 2},
+					{value: "three", score: 3}, {value: "four", score: 4},
+					{value: "five", score: 5},
+				}),
+			},
+			command: []string{"ZMPOP", "key2", "MIN"},
+			expectedValues: map[string]*SortedSet{
+				"key2": NewSortedSet([]MemberParam{
+					{value: "two", score: 2},
+					{value: "three", score: 3}, {value: "four", score: 4},
+					{value: "five", score: 5},
+				}),
+			},
+			expectedResponse: [][]string{
+				{"one", "1"},
+			},
+			expectedError: nil,
+		},
+		{ // 3. Successfully pop one max element by specifying MAX modifier
+			preset: true,
+			presetValues: map[string]interface{}{
+				"key3": NewSortedSet([]MemberParam{
+					{value: "one", score: 1}, {value: "two", score: 2},
+					{value: "three", score: 3}, {value: "four", score: 4},
+					{value: "five", score: 5},
+				}),
+			},
+			command: []string{"ZMPOP", "key3", "MAX"},
+			expectedValues: map[string]*SortedSet{
+				"key3": NewSortedSet([]MemberParam{
+					{value: "one", score: 1}, {value: "two", score: 2},
+					{value: "three", score: 3}, {value: "four", score: 4},
+				}),
+			},
+			expectedResponse: [][]string{
+				{"five", "5"},
+			},
+			expectedError: nil,
+		},
+		{ // 4. Successfully pop multiple min elements
+			preset: true,
+			presetValues: map[string]interface{}{
+				"key4": NewSortedSet([]MemberParam{
+					{value: "one", score: 1}, {value: "two", score: 2},
+					{value: "three", score: 3}, {value: "four", score: 4},
+					{value: "five", score: 5}, {value: "six", score: 6},
+				}),
+			},
+			command: []string{"ZMPOP", "key4", "MIN", "COUNT", "5"},
+			expectedValues: map[string]*SortedSet{
+				"key4": NewSortedSet([]MemberParam{
+					{value: "six", score: 6},
+				}),
+			},
+			expectedResponse: [][]string{
+				{"one", "1"}, {"two", "2"}, {"three", "3"},
+				{"four", "4"}, {"five", "5"},
+			},
+			expectedError: nil,
+		},
+		{ // 5. Successfully pop multiple max elements
+			preset: true,
+			presetValues: map[string]interface{}{
+				"key5": NewSortedSet([]MemberParam{
+					{value: "one", score: 1}, {value: "two", score: 2},
+					{value: "three", score: 3}, {value: "four", score: 4},
+					{value: "five", score: 5}, {value: "six", score: 6},
+				}),
+			},
+			command: []string{"ZMPOP", "key5", "MAX", "COUNT", "5"},
+			expectedValues: map[string]*SortedSet{
+				"key5": NewSortedSet([]MemberParam{
+					{value: "one", score: 1},
+				}),
+			},
+			expectedResponse: [][]string{{"two", "2"}, {"three", "3"}, {"four", "4"}, {"five", "5"}, {"six", "6"}},
+			expectedError:    nil,
+		},
+		{ // 6. Successfully pop elements from the first set which is non-empty
+			preset: true,
+			presetValues: map[string]interface{}{
+				"key6": NewSortedSet([]MemberParam{}),
+				"key7": NewSortedSet([]MemberParam{
+					{value: "one", score: 1}, {value: "two", score: 2},
+					{value: "three", score: 3}, {value: "four", score: 4},
+					{value: "five", score: 5}, {value: "six", score: 6},
+				}),
+			},
+			command: []string{"ZMPOP", "key6", "key7", "MAX", "COUNT", "5"},
+			expectedValues: map[string]*SortedSet{
+				"key6": NewSortedSet([]MemberParam{}),
+				"key7": NewSortedSet([]MemberParam{
+					{value: "one", score: 1},
+				}),
+			},
+			expectedResponse: [][]string{{"two", "2"}, {"three", "3"}, {"four", "4"}, {"five", "5"}, {"six", "6"}},
+			expectedError:    nil,
+		},
+		{ // 7. Skip the non-set items and pop elements from the first non-empty sorted set found
+			preset: true,
+			presetValues: map[string]interface{}{
+				"key8":  "Default value",
+				"key9":  56,
+				"key10": NewSortedSet([]MemberParam{}),
+				"key11": NewSortedSet([]MemberParam{
+					{value: "one", score: 1}, {value: "two", score: 2},
+					{value: "three", score: 3}, {value: "four", score: 4},
+					{value: "five", score: 5}, {value: "six", score: 6},
+				}),
+			},
+			command: []string{"ZMPOP", "key8", "key9", "key10", "key11", "MIN", "COUNT", "5"},
+			expectedValues: map[string]*SortedSet{
+				"key10": NewSortedSet([]MemberParam{}),
+				"key11": NewSortedSet([]MemberParam{
+					{value: "six", score: 6},
+				}),
+			},
+			expectedResponse: [][]string{{"one", "1"}, {"two", "2"}, {"three", "3"}, {"four", "4"}, {"five", "5"}},
+			expectedError:    nil,
+		},
+		{ // 9. Return error when count is a negative integer
+			preset:        false,
+			command:       []string{"ZMPOP", "key8", "MAX", "COUNT", "-20"},
+			expectedError: errors.New("count must be a positive integer"),
+		},
+		{ // 9. Command too short
+			preset:        false,
+			command:       []string{"ZMPOP"},
+			expectedError: errors.New(utils.WRONG_ARGS_RESPONSE),
+		},
+	}
+
+	for _, test := range tests {
+		if test.preset {
+			for key, value := range test.presetValues {
+				if _, err := mockServer.CreateKeyAndLock(context.Background(), key); err != nil {
+					t.Error(err)
+				}
+				mockServer.SetValue(context.Background(), key, value)
+				mockServer.KeyUnlock(key)
+			}
+		}
+		res, err := handleZMPOP(context.Background(), test.command, mockServer, nil)
+		if test.expectedError != nil {
+			if err.Error() != test.expectedError.Error() {
+				t.Errorf("expected error \"%s\", got \"%s\"", test.expectedError.Error(), err.Error())
+			}
+			continue
+		}
+		if err != nil {
+			t.Error(err)
+		}
+		rd := resp.NewReader(bytes.NewBuffer(res))
+		rv, _, err := rd.ReadValue()
+		if err != nil {
+			t.Error(err)
+		}
+		for _, element := range rv.Array() {
+			if !slices.ContainsFunc(test.expectedResponse, func(expected []string) bool {
+				return element.Array()[0].String() == expected[0] && element.Array()[1].String() == expected[1]
+			}) {
+				t.Errorf("expected response %+v, got %+v", test.expectedResponse, rv.Array())
+			}
+		}
+		for key, expectedSortedSet := range test.expectedValues {
+			if _, err = mockServer.KeyRLock(context.Background(), key); err != nil {
+				t.Error(err)
+			}
+			set, ok := mockServer.GetValue(key).(*SortedSet)
+			if !ok {
+				t.Errorf("expected key \"%s\" to be a sorted set, got another type", key)
+			}
+			if !set.Equals(expectedSortedSet) {
+				t.Errorf("expected sorted set at key \"%s\" %+v, got %+v", key, expectedSortedSet, set)
+			}
+		}
+	}
+}
+
+func Test_HandleZPOP(t *testing.T) {}
+
+func Test_HandleZMSCORE(t *testing.T) {}
+
+func Test_HandleZRANDMEMBER(t *testing.T) {}
+
+func Test_HandleZRANK(t *testing.T) {}
+
+func Test_HandleZREM(t *testing.T) {}
+
+func Test_HandleZSCORE(t *testing.T) {}
+
+func Test_HandleZREMRANGEBYSCORE(t *testing.T) {}
+
+func Test_HandleZREMRANGEBYRANK(t *testing.T) {}
+
+func Test_HandleZREMRANGEBYLEX(t *testing.T) {}
+
+func Test_HandleZRANGE(t *testing.T) {}
+
+func Test_HandleZRANGESTORE(t *testing.T) {}
+
 func Test_HandleZINTER(t *testing.T) {
 	mockServer := server.NewServer(server.Opts{})
 
@@ -1812,30 +2050,6 @@ func Test_HandleZINTERSTORE(t *testing.T) {
 	}
 }
 
-func Test_HandleZMPOP(t *testing.T) {}
-
-func Test_HandleZPOP(t *testing.T) {}
-
-func Test_HandleZMSCORE(t *testing.T) {}
-
-func Test_HandleZRANDMEMBER(t *testing.T) {}
-
-func Test_HandleZRANK(t *testing.T) {}
-
-func Test_HandleZREM(t *testing.T) {}
-
-func Test_HandleZSCORE(t *testing.T) {}
-
-func Test_HandleZREMRANGEBYSCORE(t *testing.T) {}
-
-func Test_HandleZREMRANGEBYRANK(t *testing.T) {}
-
-func Test_HandleZREMRANGEBYLEX(t *testing.T) {}
-
-func Test_HandleZRANGE(t *testing.T) {}
-
-func Test_HandleZRANGESTORE(t *testing.T) {}
-
 func Test_HandleZUNION(t *testing.T) {
 	mockServer := server.NewServer(server.Opts{})
 

src/modules/sorted_set/sorted_set.go

@@ -3,7 +3,6 @@ package sorted_set
 import (
 	"cmp"
 	"errors"
-	"fmt"
 	"github.com/echovault/echovault/src/utils"
 	"math"
 	"math/rand"
@@ -231,16 +230,15 @@ func (set *SortedSet) Pop(count int, policy string) (*SortedSet, error) {
 	})
 
 	for i := 0; i < count; i++ {
-		if i < len(members) {
+		if i >= len(members) {
+			break
+		}
 		set.Remove(members[i].value)
 		_, err := popped.AddOrUpdate([]MemberParam{members[i]}, nil, nil, nil, nil)
 		if err != nil {
-				fmt.Println(err.Error())
-				// TODO: Add all the removed elements back if we encounter an error
 			return nil, err
 		}
 	}
-	}
 
 	return popped, nil
 }
@@ -263,6 +261,24 @@ type SortedSetParam struct {
 	weight int
 }
 
+func (set *SortedSet) Equals(other *SortedSet) bool {
+	if set.Cardinality() != other.Cardinality() {
+		return false
+	}
+	if set.Cardinality() == 0 {
+		return true
+	}
+	for _, member := range set.members {
+		if !other.Contains(member.value) {
+			return false
+		}
+		if member.score != other.Get(member.value).score {
+			return false
+		}
+	}
+	return true
+}
+
 // Union uses divided & conquer to calculate the union of multiple sets
 func Union(aggregate string, setParams ...SortedSetParam) *SortedSet {
 	switch len(setParams) {