add support for ZADD INCR

Author: alicebob

cmd_sorted_set.go

@@ -55,8 +55,11 @@ func (m *Miniredis) cmdZadd(c *server.Peer, cmd string, args []string) {
 		nx    = false
 		xx    = false
 		ch    = false
+		incr  = false
 		elems = map[string]float64{}
 	)
+
+outer:
 	for len(args) > 0 {
 		switch strings.ToUpper(args[0]) {
 		case "NX":
@@ -71,21 +74,29 @@ func (m *Miniredis) cmdZadd(c *server.Peer, cmd string, args []string) {
 			ch = true
 			args = args[1:]
 			continue
+		case "INCR":
+			incr = true
+			args = args[1:]
+			continue
 		default:
-			if len(args) < 2 {
-				setDirty(c)
-				c.WriteError(msgSyntaxError)
-				return
-			}
-			score, err := strconv.ParseFloat(args[0], 64)
-			if err != nil {
-				setDirty(c)
-				c.WriteError(msgInvalidFloat)
-				return
-			}
-			elems[args[1]] = score
-			args = args[2:]
+			break outer
+		}
+	}
+
+	if len(args) == 0 || len(args)%2 != 0 {
+		setDirty(c)
+		c.WriteError(msgSyntaxError)
+		return
+	}
+	for len(args) > 0 {
+		score, err := strconv.ParseFloat(args[0], 64)
+		if err != nil {
+			setDirty(c)
+			c.WriteError(msgInvalidFloat)
+			return
 		}
+		elems[args[1]] = score
+		args = args[2:]
 	}
 
 	if xx && nx {
@@ -94,6 +105,12 @@ func (m *Miniredis) cmdZadd(c *server.Peer, cmd string, args []string) {
 		return
 	}
 
+	if incr && len(elems) > 1 {
+		setDirty(c)
+		c.WriteError(msgSingleElementPair)
+		return
+	}
+
 	withTx(m, c, func(c *server.Peer, ctx *connCtx) {
 		db := m.db(ctx.selectedDB)
 
@@ -102,6 +119,22 @@ func (m *Miniredis) cmdZadd(c *server.Peer, cmd string, args []string) {
 			return
 		}
 
+		if incr {
+			for member, delta := range elems {
+				if nx && db.ssetExists(key, member) {
+					c.WriteNull()
+					return
+				}
+				if xx && !db.ssetExists(key, member) {
+					c.WriteNull()
+					return
+				}
+				newScore := db.ssetIncrby(key, member, delta)
+				c.WriteBulk(formatFloat(newScore))
+			}
+			return
+		}
+
 		res := 0
 		for member, score := range elems {
 			if nx && db.ssetExists(key, member) {

cmd_sorted_set_test.go

@@ -172,6 +172,34 @@ func TestSortedSetAdd(t *testing.T) {
 		ok(t, err)
 		equals(t, 1, b)
 
+		n, err := redis.Float64(c.Do("ZADD", "z", "INCR", 1.2, "one"))
+		ok(t, err)
+		equals(t, 2.4, n)
+
+		z, err := c.Do("ZADD", "z", "INCR", "NX", 1.2, "one")
+		ok(t, err)
+		equals(t, nil, z)
+
+		n, err = redis.Float64(c.Do("ZADD", "z", "INCR", "XX", 1.2, "one"))
+		ok(t, err)
+		equals(t, 3.6, n)
+
+		z, err = c.Do("ZADD", "q", "INCR", "XX", 1.2, "one")
+		ok(t, err)
+		equals(t, nil, z)
+
+		n, err = redis.Float64(c.Do("ZADD", "q", "INCR", "NX", 1.2, "one"))
+		ok(t, err)
+		equals(t, 1.2, n)
+
+		z, err = c.Do("ZADD", "q", "INCR", "NX", 1.2, "one")
+		ok(t, err)
+		equals(t, nil, z)
+
+		// CH is ignored with INCR
+		n, err = redis.Float64(c.Do("ZADD", "z", "INCR", "CH", 1.2, "one"))
+		ok(t, err)
+		equals(t, 4.8, n)
 	}
 
 	// Error cases
@@ -192,8 +220,12 @@ func TestSortedSetAdd(t *testing.T) {
 		assert(t, err != nil, "ZADD error")
 		_, err = redis.String(c.Do("ZADD", "set", "MX", 1.0))
 		assert(t, err != nil, "ZADD error")
+		_, err = redis.String(c.Do("ZADD", "set", 1.0, "key", "MX"))
+		assert(t, err != nil, "ZADD error")
 		_, err = redis.String(c.Do("ZADD", "set", "MX", "XX", 1.0, "foo"))
 		assert(t, err != nil, "ZADD error")
+		_, err = redis.String(c.Do("ZADD", "set", "INCR", 1.0, "foo", 2.3, "bar"))
+		assert(t, err != nil, "ZADD error")
 	}
 }
 

redis.go

@@ -29,6 +29,7 @@ const (
 	msgInvalidKeysNumber  = "ERR Number of keys can't be greater than number of args"
 	msgNegativeKeysNumber = "ERR Number of keys can't be negative"
 	msgScriptUsage        = "ERR Unknown SCRIPT subcommand or wrong # of args."
+	msgSingleElementPair  = "ERR INCR option supports a single increment-element pair"
 	msgNoScriptFound      = "NOSCRIPT No matching script. Please use EVAL."
 )