support for unsorted sets within zinterstore/zunionstore

lib/mock_redis/zset_methods.rb

@@ -282,7 +282,7 @@ def combine_weighted_zsets(keys, options, how)
                      raise Redis::CommandError, 'ERR syntax error'
                    end
 
-      with_zsets_at(*keys) do |*zsets|
+      with_zsets_at(*keys, coercible: true) do |*zsets|
         zsets.zip(weights).map do |(zset, weight)|
           zset.reduce(Zset.new) do |acc, (score, member)|
             acc.add(score * weight, member)
@@ -293,16 +293,30 @@ def combine_weighted_zsets(keys, options, how)
       end
     end
 
-    def with_zset_at(key, &blk)
-      with_thing_at(key, :assert_zsety, proc { Zset.new }, &blk)
+    def coerce_to_zset(set)
+      zset = Zset.new
+      set.each do |member|
+        zset.add(1.0, member)
+      end
+      zset
     end
 
-    def with_zsets_at(*keys, &blk)
+    def with_zset_at(key, coercible: false, &blk)
+      if coercible
+        with_thing_at(key, :assert_coercible_zsety, proc { Zset.new }) do |value|
+          blk.call value.is_a?(Set) ? coerce_to_zset(value) : value
+        end
+      else
+        with_thing_at(key, :assert_zsety, proc { Zset.new }, &blk)
+      end
+    end
+
+    def with_zsets_at(*keys, coercible: false, &blk)
       if keys.length == 1
-        with_zset_at(keys.first, &blk)
+        with_zset_at(keys.first, coercible: coercible, &blk)
       else
-        with_zset_at(keys.first) do |set|
-          with_zsets_at(*(keys[1..-1])) do |*sets|
+        with_zset_at(keys.first, coercible: coercible) do |set|
+          with_zsets_at(*(keys[1..-1]), coercible: coercible) do |*sets|
             yield(*([set] + sets))
           end
         end
@@ -313,13 +327,24 @@ def zsety?(key)
       data[key].nil? || data[key].is_a?(Zset)
     end
 
+    def coercible_zsety?(key)
+      zsety?(key) || data[key].is_a?(Set)
+    end
+
     def assert_zsety(key)
       unless zsety?(key)
         raise Redis::CommandError,
           'WRONGTYPE Operation against a key holding the wrong kind of value'
       end
     end
 
+    def assert_coercible_zsety(key)
+      unless coercible_zsety?(key)
+        raise Redis::CommandError,
+          'WRONGTYPE Operation against a key holding the wrong kind of value'
+      end
+    end
+
     def looks_like_float?(x)
       # ugh, exceptions for flow control.
       !!Float(x) rescue false

spec/commands/zinterstore_spec.rb

@@ -42,6 +42,40 @@
     end.should raise_error(Redis::CommandError)
   end
 
+  context 'when used with a set' do
+    before do
+      @primes_text = 'mock-redis-test:zinterstore:primes-text'
+
+      @redises.sadd(@primes_text, 'two')
+      @redises.sadd(@primes_text, 'three')
+      @redises.sadd(@primes_text, 'five')
+      @redises.sadd(@primes_text, 'seven')
+    end
+
+    it 'returns the number of elements in the new set' do
+      @redises.zinterstore(@dest, [@odds, @primes_text]).should == 3
+    end
+
+    it 'sums the scores, substituting 1.0 for set values' do
+      @redises.zinterstore(@dest, [@odds, @primes_text])
+      @redises.zrange(@dest, 0, -1, :with_scores => true).should ==
+        [['three', 4.0], ['five', 6.0], ['seven', 8.0]]
+    end
+  end
+
+  context 'when used with a non-coercible structure' do
+    before do
+      @non_set = 'mock-redis-test:zinterstore:non-set'
+
+      @redises.set(@non_set, 'one')
+    end
+    it 'raises an error for wrong value type' do
+      lambda do
+        @redises.zinterstore(@dest, [@odds, @non_set])
+      end.should raise_error(Redis::CommandError)
+    end
+  end
+
   context 'the :weights argument' do
     it 'multiplies the scores by the weights while aggregating' do
       @redises.zinterstore(@dest, [@odds, @primes], :weights => [2, 3])

spec/commands/zunionstore_spec.rb

@@ -41,6 +41,39 @@
     end.should raise_error(Redis::CommandError)
   end
 
+  context 'when used with a set' do
+    before do
+      @set4 = 'mock-redis-test:zunionstore4'
+
+      @redises.sadd(@set4, 'two')
+      @redises.sadd(@set4, 'three')
+      @redises.sadd(@set4, 'four')
+    end
+
+    it 'returns the number of elements in the new set' do
+      @redises.zunionstore(@dest, [@set3, @set4]).should == 4
+    end
+
+    it 'sums the scores, substituting 1.0 for set values' do
+      @redises.zunionstore(@dest, [@set3, @set4])
+      @redises.zrange(@dest, 0, -1, :with_scores => true).should ==
+        [['four', 1.0], ['one', 1.0], ['two', 3.0], ['three', 4.0]]
+    end
+  end
+
+  context 'when used with a non-coercible structure' do
+    before do
+      @non_set = 'mock-redis-test:zunionstore4'
+
+      @redises.set(@non_set, 'one')
+    end
+    it 'raises an error for wrong value type' do
+      lambda do
+        @redises.zunionstore(@dest, [@set1, @non_set])
+      end.should raise_error(Redis::CommandError)
+    end
+  end
+
   context 'the :weights argument' do
     it 'multiplies the scores by the weights while aggregating' do
       @redises.zunionstore(@dest, [@set1, @set2, @set3], :weights => [2, 3, 5])