Implemented test case for ZADD command handler. Added error returns for incompatible flags in ZADD commands

src/modules/sorted_set/commands.go

@@ -94,10 +94,19 @@ func handleZADD(ctx context.Context, cmd []string, server utils.Server, conn *ne
 		for _, option := range options {
 			if slices.Contains([]string{"xx", "nx"}, strings.ToLower(option)) {
 				updatePolicy = option
+				// If option is "NX" and comparison is not nil, return an error
+				if strings.EqualFold(option, "NX") && comparison != nil {
+					return nil, errors.New("GT/LT flags not allowed if NX flag is provided")
+				}
 				continue
 			}
 			if slices.Contains([]string{"gt", "lt"}, strings.ToLower(option)) {
 				comparison = option
+				// If updatePolicy is "NX", return an error
+				up, _ := updatePolicy.(string)
+				if strings.EqualFold(up, "NX") {
+					return nil, errors.New("GT/LT flags not allowed if NX flag is provided")
+				}
 				continue
 			}
 			if strings.EqualFold(option, "ch") {
@@ -106,6 +115,10 @@ func handleZADD(ctx context.Context, cmd []string, server utils.Server, conn *ne
 			}
 			if strings.EqualFold(option, "incr") {
 				incr = option
+				// If members length is more than 1, return an error
+				if len(members) > 1 {
+					return nil, errors.New("cannot pass more than one score/member pair when INCR flag is provided")
+				}
 				continue
 			}
 			return nil, fmt.Errorf("invalid option %s", option)
@@ -137,8 +150,7 @@ func handleZADD(ctx context.Context, cmd []string, server utils.Server, conn *ne
 	}
 
 	// Key does not exist
-	_, err := server.CreateKeyAndLock(ctx, key)
+	if _, err := server.CreateKeyAndLock(ctx, key); err != nil {
-	if err != nil {
 		return nil, err
 	}
 	defer server.KeyUnlock(key)
@@ -1552,7 +1564,13 @@ func Commands() []utils.Command {
 			Command:    "zadd",
 			Categories: []string{utils.SortedSetCategory, utils.WriteCategory, utils.FastCategory},
 			Description: `(ZADD key [NX | XX] [GT | LT] [CH] [INCR] score member [score member...])
-Adds all the specified members with the specified scores to the sorted set at the key`,
+Adds all the specified members with the specified scores to the sorted set at the key.
+"NX" only adds the member if it currently does not exist in the sorted set.
+"XX" only updates the scores of members that exist in the sorted set.
+"GT"" only updates the score if the new score is greater than the current score.
+"LT" only updates the score if the new score is less than the current score.
+"CH" modifies the result to return total number of members changed + added, instead of only new members added.
+"INCR" modifies the command to act like ZINCRBY, only one score/member pair can be specified in this mode.`,
 			Sync:              true,
 			KeyExtractionFunc: zaddKeyFunc,
 			HandlerFunc:       handleZADD,

src/modules/sorted_set/commands_test.go

@@ -1,10 +1,267 @@
 package sorted_set
 
 import (
+	"bytes"
+	"context"
+	"errors"
+	"github.com/echovault/echovault/src/server"
+	"github.com/echovault/echovault/src/utils"
+	"github.com/tidwall/resp"
 	"testing"
 )
 
-func Test_HandleZADD(t *testing.T) {}
+func Test_HandleZADD(t *testing.T) {
+	mockServer := server.NewServer(server.Opts{})
+
+	tests := []struct {
+		preset           bool
+		presetValue      *SortedSet
+		key              string
+		command          []string
+		expectedValue    *SortedSet
+		expectedResponse int
+		expectedError    error
+	}{
+		{ // 1. Create new sorted set and return the cardinality of the new sorted set.
+			preset:      false,
+			presetValue: nil,
+			key:         "key1",
+			command:     []string{"ZADD", "key1", "5.5", "member1", "67.77", "member2", "10", "member3"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			expectedResponse: 3,
+			expectedError:    nil,
+		},
+		{ // 2. Only add the elements that do not currently exist in the sorted set when NX flag is provided
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			key:     "key2",
+			command: []string{"ZADD", "key2", "NX", "5.5", "member1", "67.77", "member4", "10", "member5"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+				{value: "member4", score: Score(67.77)},
+				{value: "member5", score: Score(10)},
+			}),
+			expectedResponse: 2,
+			expectedError:    nil,
+		},
+		{ // 3. Do not add any elements when providing existing members with NX flag
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			key:     "key3",
+			command: []string{"ZADD", "key3", "NX", "5.5", "member1", "67.77", "member2", "10", "member3"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			expectedResponse: 0,
+			expectedError:    nil,
+		},
+		{ // 4. Successfully add elements to an existing set when XX flag is provided with existing elements
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			key:     "key4",
+			command: []string{"ZADD", "key4", "XX", "CH", "55", "member1", "1005", "member2", "15", "member3", "99.75", "member4"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(55)},
+				{value: "member2", score: Score(1005)},
+				{value: "member3", score: Score(15)},
+			}),
+			expectedResponse: 3,
+			expectedError:    nil,
+		},
+		{ // 5. Fail to add element when providing XX flag with elements that do not exist in the sorted set.
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			key:     "key5",
+			command: []string{"ZADD", "key5", "XX", "5.5", "member4", "100.5", "member5", "15", "member6"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			expectedResponse: 0,
+			expectedError:    nil,
+		},
+		{ // 6. Only update the elements where provided score is greater than current score if GT flag
+			// Return only the new elements added by default
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			key:     "key6",
+			command: []string{"ZADD", "key6", "XX", "CH", "GT", "7.5", "member1", "100.5", "member4", "15", "member5"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(7.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			expectedResponse: 1,
+			expectedError:    nil,
+		},
+		{ // 7. Only update the elements where provided score is less than current score if LT flag is provided
+			// Return only the new elements added by default.
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			key:     "key7",
+			command: []string{"ZADD", "key7", "XX", "LT", "3.5", "member1", "100.5", "member4", "15", "member5"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(3.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			expectedResponse: 0,
+			expectedError:    nil,
+		},
+		{ // 8. Return all the elements that were updated AND added when CH flag is provided
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			key:     "key8",
+			command: []string{"ZADD", "key8", "XX", "LT", "CH", "3.5", "member1", "100.5", "member4", "15", "member5"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(3.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			expectedResponse: 1,
+			expectedError:    nil,
+		},
+		{ // 9. Increment the member by score
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(10)},
+			}),
+			key:     "key9",
+			command: []string{"ZADD", "key9", "INCR", "5.5", "member3"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "member1", score: Score(5.5)},
+				{value: "member2", score: Score(67.77)},
+				{value: "member3", score: Score(15.5)},
+			}),
+			expectedResponse: 0,
+			expectedError:    nil,
+		},
+		{ // 10. Fail when GT/LT flag is provided alongside NX flag
+			preset:           false,
+			presetValue:      nil,
+			key:              "key10",
+			command:          []string{"ZADD", "key10", "NX", "LT", "CH", "3.5", "member1", "100.5", "member4", "15", "member5"},
+			expectedValue:    nil,
+			expectedResponse: 0,
+			expectedError:    errors.New("GT/LT flags not allowed if NX flag is provided"),
+		},
+		{ // 11. Command is too short
+			preset:           false,
+			presetValue:      nil,
+			key:              "key11",
+			command:          []string{"ZADD", "key11"},
+			expectedValue:    nil,
+			expectedResponse: 0,
+			expectedError:    errors.New(utils.WRONG_ARGS_RESPONSE),
+		},
+		{ // 12. Throw error when score/member entries are do not match
+			preset:           false,
+			presetValue:      nil,
+			key:              "key11",
+			command:          []string{"ZADD", "key12", "10.5", "member1", "12.5"},
+			expectedValue:    nil,
+			expectedResponse: 0,
+			expectedError:    errors.New("score/member pairs must be float/string"),
+		},
+		{ // 13. Throw error when INCR flag is passed with more than one score/member pair
+			preset:           false,
+			presetValue:      nil,
+			key:              "key13",
+			command:          []string{"ZADD", "key13", "INCR", "10.5", "member1", "12.5", "member2"},
+			expectedValue:    nil,
+			expectedResponse: 0,
+			expectedError:    errors.New("cannot pass more than one score/member pair when INCR flag is provided"),
+		},
+	}
+
+	for _, test := range tests {
+		if test.preset {
+			if _, err := mockServer.CreateKeyAndLock(context.Background(), test.key); err != nil {
+				t.Error(err)
+			}
+			mockServer.SetValue(context.Background(), test.key, test.presetValue)
+			mockServer.KeyUnlock(test.key)
+		}
+		res, err := handleZADD(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.NewReader(res))
+		rv, _, err := rd.ReadValue()
+		if err != nil {
+			t.Error(err)
+		}
+		if rv.Integer() != test.expectedResponse {
+			t.Errorf("expected response %d at key \"%s\", got %d", test.expectedResponse, test.key, rv.Integer())
+		}
+		// Fetch the sorted set from the server and check it against the expected result
+		if _, err = mockServer.KeyRLock(context.Background(), test.key); err != nil {
+			t.Error(err)
+		}
+		sortedSet, ok := mockServer.GetValue(test.key).(*SortedSet)
+		if !ok {
+			t.Errorf("expected the value at key \"%s\" to be a sorted set, got another type", test.key)
+		}
+		if test.expectedValue == nil {
+			continue
+		}
+		for _, member := range sortedSet.GetAll() {
+			expectedMember := test.expectedValue.Get(member.value)
+			if !expectedMember.exists {
+				t.Errorf("could not find member %+v in expected sorted set, found in stored set", member)
+			}
+			if member.score != expectedMember.score {
+				t.Errorf("expected member \"%s\" to have score %f, got score %f", expectedMember.value, expectedMember.score, member.score)
+			}
+		}
+		mockServer.KeyRUnlock(test.key)
+	}
+}
 
 func Test_HandleZCARD(t *testing.T) {}