Changed count to options, fixed small parts of the code

Signed-off-by: Edward Liang <[email protected]>

go/api/base_client.go

@@ -4535,8 +4535,10 @@ func (client *baseClient) BZPopMin(keys []string, timeoutSecs float64) (Result[K
 // Parameters:
 //
 //	keys          - An array of keys to lists.
-//	scoreFilter   - The element pop criteria - either [api.MIN] or [api.MAX] to pop members with the lowest/highest scores accordingly.
-//	timeoutSecs   - The number of seconds to wait for a blocking operation to complete. A value of `0` will block indefinitely.
+//	scoreFilter   - The element pop criteria - either [api.MIN] or [api.MAX] to pop members with the lowest/highest
+//					scores accordingly.
+//	timeoutSecs   - The number of seconds to wait for a blocking operation to complete. A value of `0` will block
+//					indefinitely.
 //
 // Return value:
 //
@@ -4599,7 +4601,8 @@ func (client *baseClient) BZMPop(
 // Parameters:
 //
 //	keys          - An array of keys to lists.
-//	scoreFilter   - The element pop criteria - either [api.MIN] or [api.MAX] to pop members with the lowest/highest scores accordingly.
+//	scoreFilter   - The element pop criteria - either [api.MIN] or [api.MAX] to pop members with the lowest/highest
+//					scores accordingly.
 //	count         - The maximum number of popped elements.
 //	timeoutSecs   - The number of seconds to wait for a blocking operation to complete. A value of `0` will block indefinitely.
 //
@@ -4613,16 +4616,16 @@ func (client *baseClient) BZMPop(
 // For example:
 //
 //	result, err := client.ZAdd("my_list", map[string]float64{"five": 5.0, "six": 6.0})
-//	result, err := client.BZMPopCount([]string{"my_list"}, api.MAX, 2, 0.1)
+//	result, err := client.BZMPopWithOptions([]string{"my_list"}, api.MAX, 0.1, options.NewZMPopOptions().SetCount(2))
 //	result["my_list"] = []MemberAndScore{{Member: "six", Score: 6.0}, {Member: "five", Score 5.0}}
 //
 // [valkey.io]: https://valkey.io/commands/bzmpop/
 // [Blocking Commands]: https://github.com/valkey-io/valkey-glide/wiki/General-Concepts#blocking-commands
-func (client *baseClient) BZMPopCount(
+func (client *baseClient) BZMPopWithOptions(
 	keys []string,
 	scoreFilter ScoreFilter,
-	count int64,
 	timeoutSecs float64,
+	opts *options.ZMPopOptions,
 ) (Result[KeyWithArrayOfMembersAndScores], error) {
 	scoreFilterStr, err := scoreFilter.toString()
 	if err != nil {
@@ -4640,7 +4643,14 @@ func (client *baseClient) BZMPopCount(
 	args := make([]string, 0, len(keys)+5)
 	args = append(args, utils.FloatToString(timeoutSecs), strconv.Itoa(len(keys)))
 	args = append(args, keys...)
-	args = append(args, scoreFilterStr, CountKeyword, utils.IntToString(count))
+	args = append(args, scoreFilterStr)
+	if opts != nil {
+		optionArgs, err := opts.ToArgs()
+		if err != nil {
+			return CreateNilKeyWithArrayOfMembersAndScoresResult(), err
+		}
+		args = append(args, optionArgs...)
+	}
 	result, err := client.executeCommand(C.BZMPop, args)
 	if err != nil {
 		return CreateNilKeyWithArrayOfMembersAndScoresResult(), err

go/api/options/zmpop_options.go

@@ -0,0 +1,34 @@
+// Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
+
+package options
+
+import (
+	"github.com/valkey-io/valkey-glide/go/glide/utils"
+)
+
+// Optional arguments for `ZMPop` and `BZMPop` in [SortedSetCommands]
+type ZMPopOptions struct {
+	count      int64
+	countIsSet bool
+}
+
+func NewZMPopOptions() *ZMPopOptions {
+	return &ZMPopOptions{}
+}
+
+// Set the count.
+func (zmpo *ZMPopOptions) SetCount(count int64) *ZMPopOptions {
+	zmpo.count = count
+	zmpo.countIsSet = true
+	return zmpo
+}
+
+func (zmpo *ZMPopOptions) ToArgs() ([]string, error) {
+	var args []string
+
+	if zmpo.countIsSet {
+		args = append(args, "COUNT", utils.IntToString(zmpo.count))
+	}
+
+	return args, nil
+}

go/api/response_handlers.go

@@ -542,7 +542,7 @@ func handleKeyWithArrayOfMembersAndScoresResponse(
 ) (Result[KeyWithArrayOfMembersAndScores], error) {
 	defer C.free_command_response(response)
 
-	if response == nil || response.response_type == uint32(C.Null) {
+	if response.response_type == uint32(C.Null) {
 		return CreateNilKeyWithArrayOfMembersAndScoresResult(), nil
 	}
 
@@ -572,12 +572,10 @@ func handleKeyWithArrayOfMembersAndScoresResponse(
 			Msg: fmt.Sprintf("unexpected type of second element: %T", converted),
 		}
 	}
-	memberAndScoreArray := make([]MemberAndScore, len(res))
+	memberAndScoreArray := make([]MemberAndScore, 0, len(res))
 
-	idx := 0
 	for k, v := range res {
-		memberAndScoreArray[idx] = MemberAndScore{k, v}
-		idx++
+		memberAndScoreArray = append(memberAndScoreArray, MemberAndScore{k, v})
 	}
 
 	// Ensure consistent output

go/api/sorted_set_commands.go

@@ -38,11 +38,11 @@ type SortedSetCommands interface {
 
 	BZMPop(keys []string, scoreFilter ScoreFilter, timeoutSecs float64) (Result[KeyWithArrayOfMembersAndScores], error)
 
-	BZMPopCount(
+	BZMPopWithOptions(
 		keys []string,
 		scoreFilter ScoreFilter,
-		count int64,
 		timeoutSecs float64,
+		options *options.ZMPopOptions,
 	) (Result[KeyWithArrayOfMembersAndScores], error)
 
 	ZRange(key string, rangeQuery options.ZRangeQuery) ([]string, error)

go/integTest/shared_commands_test.go

@@ -2886,7 +2886,7 @@ func (suite *GlideTestSuite) TestBLMPopAndBLMPopCount() {
 	})
 }
 
-func (suite *GlideTestSuite) TestBZMPopAndBZMPopCount() {
+func (suite *GlideTestSuite) TestBZMPopAndBZMPopWithOptions() {
 	if suite.serverVersion < "7.0.0" {
 		suite.T().Skip("This feature is added in version 7")
 	}
@@ -2913,7 +2913,7 @@ func (suite *GlideTestSuite) TestBZMPopAndBZMPopCount() {
 		assert.Equal(suite.T(), int64(3), res4)
 
 		// Try to pop the top 2 elements from key1
-		res5, err := client.BZMPopCount([]string{key1}, api.MAX, int64(2), float64(0.1))
+		res5, err := client.BZMPopWithOptions([]string{key1}, api.MAX, float64(0.1), options.NewZMPopOptions().SetCount(2))
 		assert.Nil(suite.T(), err)
 		assert.Equal(
 			suite.T(),