Merge pull request #141 from alicebob/setstore

use a normal set as a ZUNIONSTORE source

Author: alicebob

cmd_sorted_set.go

@@ -1265,32 +1265,42 @@ func (m *Miniredis) cmdZunionstore(c *server.Peer, cmd string, args []string) {
 			if !db.exists(key) {
 				continue
 			}
-			if db.t(key) != "zset" {
+
+			var set map[string]float64
+			switch db.t(key) {
+			case "set":
+				set = map[string]float64{}
+				for elem := range db.setKeys[key] {
+					set[elem] = 1.0
+				}
+			case "zset":
+				set = db.sortedSet(key)
+			default:
 				c.WriteError(msgWrongType)
 				return
 			}
-			for _, el := range db.ssetElements(key) {
-				score := el.score
+
+			for member, score := range set {
 				if withWeights {
 					score *= weights[i]
 				}
-				old, ok := sset[el.member]
+				old, ok := sset[member]
 				if !ok {
-					sset[el.member] = score
+					sset[member] = score
 					continue
 				}
 				switch aggregate {
 				default:
 					panic("Invalid aggregate")
 				case "sum":
-					sset[el.member] += score
+					sset[member] += score
 				case "min":
 					if score < old {
-						sset[el.member] = score
+						sset[member] = score
 					}
 				case "max":
 					if score > old {
-						sset[el.member] = score
+						sset[member] = score
 					}
 				}
 			}

cmd_sorted_set_test.go

@@ -1174,19 +1174,17 @@ func TestZunionstore(t *testing.T) {
 	s.ZAdd("h2", 1.0, "field1")
 	s.ZAdd("h2", 2.0, "field2")
 
-	// Simple case
-	{
+	t.Run("simple case", func(t *testing.T) {
 		res, err := redis.Int(c.Do("ZUNIONSTORE", "new", 2, "h1", "h2"))
 		ok(t, err)
 		equals(t, 2, res)
 
 		ss, err := s.SortedSet("new")
 		ok(t, err)
 		equals(t, map[string]float64{"field1": 2, "field2": 4}, ss)
-	}
+	})
 
-	// Merge destination with itself.
-	{
+	t.Run("merge destination with itself", func(t *testing.T) {
 		s.ZAdd("h3", 1.0, "field1")
 		s.ZAdd("h3", 3.0, "field3")
 
@@ -1197,32 +1195,37 @@ func TestZunionstore(t *testing.T) {
 		ss, err := s.SortedSet("h3")
 		ok(t, err)
 		equals(t, map[string]float64{"field1": 2, "field2": 2, "field3": 3}, ss)
-	}
+	})
 
-	// WEIGHTS
-	{
+	t.Run("WEIGHTS", func(t *testing.T) {
 		res, err := redis.Int(c.Do("ZUNIONSTORE", "weighted", 2, "h1", "h2", "WeIgHtS", "4.5", "12"))
 		ok(t, err)
 		equals(t, 2, res)
 
 		ss, err := s.SortedSet("weighted")
 		ok(t, err)
 		equals(t, map[string]float64{"field1": 16.5, "field2": 33}, ss)
-	}
+	})
 
-	// AGGREGATE
-	{
+	t.Run("AGGREGATE", func(t *testing.T) {
 		res, err := redis.Int(c.Do("ZUNIONSTORE", "aggr", 2, "h1", "h2", "AgGrEgAtE", "min"))
 		ok(t, err)
 		equals(t, 2, res)
 
 		ss, err := s.SortedSet("aggr")
 		ok(t, err)
 		equals(t, map[string]float64{"field1": 1.0, "field2": 2.0}, ss)
-	}
+	})
 
-	// Wrong usage
-	{
+	t.Run("normal set", func(t *testing.T) {
+		_, err := c.Do("SADD", "set", "aap", "noot", "mies")
+		ok(t, err)
+		res, err := redis.Int(c.Do("ZUNIONSTORE", "aggr", 1, "set"))
+		ok(t, err)
+		equals(t, 3, res)
+	})
+
+	t.Run("wrong usage", func(t *testing.T) {
 		_, err := redis.Int(c.Do("ZUNIONSTORE"))
 		assert(t, err != nil, "do ZUNIONSTORE error")
 		_, err = redis.Int(c.Do("ZUNIONSTORE", "set"))
@@ -1255,7 +1258,7 @@ func TestZunionstore(t *testing.T) {
 		s.Set("str", "value")
 		_, err = redis.Int(c.Do("ZUNIONSTORE", "set", 1, "str"))
 		assert(t, err != nil, "do ZUNIONSTORE error")
-	}
+	})
 }
 
 func TestZinterstore(t *testing.T) {

integration/sorted_set_test.go

@@ -624,6 +624,13 @@ func TestZunionstore(t *testing.T) {
 		succ("TYPE", "h1"),
 		succ("TYPE", "h2"),
 	)
+	// not a sorted set, still fine
+	testCommands(t,
+		succ("SADD", "super", "1", "2", "3"),
+		succ("SADD", "exclude", "3"),
+		succ("ZUNIONSTORE", "tmp", "2", "super", "exclude", "weights", "1", "0", "aggregate", "min"),
+		succ("ZRANGE", "tmp", "0", "-1", "withscores"),
+	)
 }
 
 func TestZinterstore(t *testing.T) {