Implemented test case for ZINCRYBY command handler

src/modules/sorted_set/commands.go

@@ -469,7 +469,17 @@ func handleZINCRBY(ctx context.Context, cmd []string, server utils.Server, conn
 		increment = Score(s)
 	}
 
-	if server.KeyExists(key) {
+	if !server.KeyExists(key) {
+		// If the key does not exist, create a new sorted set at the key with
+		// the member and increment as the first value
+		if _, err = server.CreateKeyAndLock(ctx, key); err != nil {
+			return nil, err
+		}
+		server.SetValue(ctx, key, NewSortedSet([]MemberParam{{value: member, score: increment}}))
+		server.KeyUnlock(key)
+		return []byte(fmt.Sprintf("+%s\r\n\r\n", strconv.FormatFloat(float64(increment), 'f', -1, 64))), nil
+	}
+
 	if _, err = server.KeyLock(ctx, key); err != nil {
 		return nil, err
 	}
@@ -478,28 +488,17 @@ func handleZINCRBY(ctx context.Context, cmd []string, server utils.Server, conn
 	if !ok {
 		return nil, fmt.Errorf("value at %s is not a sorted set", key)
 	}
-		_, err = set.AddOrUpdate(
+	if _, err = set.AddOrUpdate(
-			[]MemberParam{{value: member, score: increment}}, "xx", nil, nil, "incr")
+		[]MemberParam{
-		if err != nil {
+			{value: member, score: increment}},
+		"xx",
+		nil,
+		nil,
+		"incr"); err != nil {
 		return nil, err
 	}
-		return []byte(fmt.Sprintf("+%f\r\n\r\n", set.Get(member).score)), nil
+	return []byte(fmt.Sprintf("+%s\r\n\r\n",
-	}
+		strconv.FormatFloat(float64(set.Get(member).score), 'f', -1, 64))), nil
-
-	if _, err = server.CreateKeyAndLock(ctx, key); err != nil {
-		return nil, err
-	}
-	defer server.KeyUnlock(key)
-
-	set := NewSortedSet([]MemberParam{
-		{
-			value: member,
-			score: increment,
-		},
-	})
-	server.SetValue(ctx, key, set)
-
-	return []byte(fmt.Sprintf("+%f\r\n\r\n", set.Get(member).score)), nil
 }
 
 func handleZINTER(ctx context.Context, cmd []string, server utils.Server, conn *net.Conn) ([]byte, error) {
@@ -1533,8 +1532,6 @@ func handleZUNIONSTORE(ctx context.Context, cmd []string, server utils.Server, c
 
 	server.SetValue(ctx, destination, union)
 
-	fmt.Println("DESTINATION: ", destination, ", CARD: ", union.Cardinality())
-
 	return []byte(fmt.Sprintf(":%d\r\n\r\n", union.Cardinality())), nil
 }
 

src/modules/sorted_set/commands_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/tidwall/resp"
 	"math"
 	"slices"
+	"strconv"
 	"testing"
 )
 
@@ -940,7 +941,226 @@ func Test_HandleZDIFFSTORE(t *testing.T) {
 	}
 }
 
-func Test_HandleZINCRBY(t *testing.T) {}
+func Test_HandleZINCRBY(t *testing.T) {
+	mockServer := server.NewServer(server.Opts{})
+
+	tests := []struct {
+		preset           bool
+		presetValue      interface{}
+		key              string
+		command          []string
+		expectedValue    *SortedSet
+		expectedResponse string
+		expectedError    error
+	}{
+		{ // 1. Successfully increment by int. Return the new score
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 1}, {value: "two", score: 2},
+				{value: "three", score: 3}, {value: "four", score: 4},
+				{value: "five", score: 5},
+			}),
+			key:     "key1",
+			command: []string{"ZINCRBY", "key1", "5", "one"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 6}, {value: "two", score: 2},
+				{value: "three", score: 3}, {value: "four", score: 4},
+				{value: "five", score: 5},
+			}),
+			expectedResponse: "6",
+			expectedError:    nil,
+		},
+		{ // 2. Successfully increment by float. Return new score
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 1}, {value: "two", score: 2},
+				{value: "three", score: 3}, {value: "four", score: 4},
+				{value: "five", score: 5},
+			}),
+			key:     "key2",
+			command: []string{"ZINCRBY", "key2", "346.785", "one"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 347.785}, {value: "two", score: 2},
+				{value: "three", score: 3}, {value: "four", score: 4},
+				{value: "five", score: 5},
+			}),
+			expectedResponse: "347.785",
+			expectedError:    nil,
+		},
+		{ // 3. Increment on non-existent sorted set will create the set with the member and increment as its score
+			preset:      false,
+			presetValue: nil,
+			key:         "key3",
+			command:     []string{"ZINCRBY", "key3", "346.785", "one"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 346.785},
+			}),
+			expectedResponse: "346.785",
+			expectedError:    nil,
+		},
+		{ // 4. Increment score to +inf
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 1}, {value: "two", score: 2},
+				{value: "three", score: 3}, {value: "four", score: 4},
+				{value: "five", score: 5},
+			}),
+			key:     "key4",
+			command: []string{"ZINCRBY", "key4", "+inf", "one"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "one", score: Score(math.Inf(1))}, {value: "two", score: 2},
+				{value: "three", score: 3}, {value: "four", score: 4},
+				{value: "five", score: 5},
+			}),
+			expectedResponse: "+Inf",
+			expectedError:    nil,
+		},
+		{ // 5. Increment score to -inf
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 1}, {value: "two", score: 2},
+				{value: "three", score: 3}, {value: "four", score: 4},
+				{value: "five", score: 5},
+			}),
+			key:     "key5",
+			command: []string{"ZINCRBY", "key5", "-inf", "one"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "one", score: Score(math.Inf(-1))}, {value: "two", score: 2},
+				{value: "three", score: 3}, {value: "four", score: 4},
+				{value: "five", score: 5},
+			}),
+			expectedResponse: "-Inf",
+			expectedError:    nil,
+		},
+		{ // 6. Incrementing score by negative increment should lower the score
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 1}, {value: "two", score: 2},
+				{value: "three", score: 3}, {value: "four", score: 4},
+				{value: "five", score: 5},
+			}),
+			key:     "key6",
+			command: []string{"ZINCRBY", "key6", "-2.5", "five"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 1}, {value: "two", score: 2},
+				{value: "three", score: 3}, {value: "four", score: 4},
+				{value: "five", score: 2.5},
+			}),
+			expectedResponse: "2.5",
+			expectedError:    nil,
+		},
+		{ // 7. Return error when attempting to increment on a value that is not a valid sorted set
+			preset:           true,
+			presetValue:      "Default value",
+			key:              "key7",
+			command:          []string{"ZINCRBY", "key7", "-2.5", "five"},
+			expectedValue:    nil,
+			expectedResponse: "",
+			expectedError:    errors.New("value at key7 is not a sorted set"),
+		},
+		{ // 8. Return error when trying to increment a member that already has score -inf
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "one", score: Score(math.Inf(-1))},
+			}),
+			key:     "key8",
+			command: []string{"ZINCRBY", "key8", "2.5", "one"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "one", score: Score(math.Inf(-1))},
+			}),
+			expectedResponse: "",
+			expectedError:    errors.New("cannot increment -inf or +inf"),
+		},
+		{ // 9. Return error when trying to increment a member that already has score +inf
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "one", score: Score(math.Inf(1))},
+			}),
+			key:     "key9",
+			command: []string{"ZINCRBY", "key9", "2.5", "one"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "one", score: Score(math.Inf(-1))},
+			}),
+			expectedResponse: "",
+			expectedError:    errors.New("cannot increment -inf or +inf"),
+		},
+		{ // 10. Return error when increment is not a valid number
+			preset: true,
+			presetValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 1},
+			}),
+			key:     "key10",
+			command: []string{"ZINCRBY", "key10", "increment", "one"},
+			expectedValue: NewSortedSet([]MemberParam{
+				{value: "one", score: 1},
+			}),
+			expectedResponse: "",
+			expectedError:    errors.New("increment must be a double"),
+		},
+		{ // 11. Command too short
+			key:              "key11",
+			command:          []string{"ZINCRBY", "key11", "one"},
+			expectedResponse: "",
+			expectedError:    errors.New(utils.WRONG_ARGS_RESPONSE),
+		},
+		{ // 12. Command too long
+			key:              "key12",
+			command:          []string{"ZINCRBY", "key12", "one", "1", "2"},
+			expectedResponse: "",
+			expectedError:    errors.New(utils.WRONG_ARGS_RESPONSE),
+		},
+	}
+
+	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 := handleZINCRBY(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)
+		}
+		if rv.String() != test.expectedResponse {
+			t.Errorf("expected response integer %s, got %s", test.expectedResponse, rv.String())
+		}
+		if test.expectedValue != nil {
+			if _, err = mockServer.KeyRLock(context.Background(), test.key); err != nil {
+				t.Error(err)
+			}
+			set, ok := mockServer.GetValue(test.key).(*SortedSet)
+			if !ok {
+				t.Errorf("expected vaule at key %s to be set, got another type", test.key)
+			}
+			for _, elem := range set.GetAll() {
+				if !test.expectedValue.Contains(elem.value) {
+					t.Errorf("could not find element %s in the expected values", elem.value)
+				}
+				if test.expectedValue.Get(elem.value).score != elem.score {
+					t.Errorf("expected score of element \"%s\" from set at key \"%s\" to be %s, got %s",
+						elem.value, test.key,
+						strconv.FormatFloat(float64(test.expectedValue.Get(elem.value).score), 'f', -1, 64),
+						strconv.FormatFloat(float64(elem.score), 'f', -1, 64),
+					)
+				}
+			}
+			mockServer.KeyRUnlock(test.key)
+		}
+	}
+}
 
 func Test_HandleZINTER(t *testing.T) {
 	mockServer := server.NewServer(server.Opts{})

src/modules/sorted_set/sorted_set.go

@@ -136,7 +136,15 @@ func (set *SortedSet) AddOrUpdate(
 	if strings.EqualFold(inc, "incr") {
 		for _, m := range members {
 			if !set.Contains(m.value) {
-				return count, fmt.Errorf("cannot increment member %s as it does not exist in the sorted set", m.value)
+				// If the member is not contained, add it with the increment as its score
+				set.members[m.value] = MemberObject{
+					value:  m.value,
+					score:  m.score,
+					exists: true,
+				}
+				// Always add count because this is the addition of a new element
+				count += 1
+				return count, err
 			}
 			if slices.Contains([]Score{Score(math.Inf(-1)), Score(math.Inf(1))}, set.members[m.value].score) {
 				return count, errors.New("cannot increment -inf or +inf")