lib/redis/connection/memory.rb in fakeredis-0.2.2 vs lib/redis/connection/memory.rb in fakeredis-0.3.0

@@ -1,5 +1,6 @@
+require 'set'
 require 'redis/connection/registry'
 require 'redis/connection/command_helper'
 
 class Redis
   module Connection
@@ -57,10 +58,13 @@
             end
           end
         end
       end
 
+      class ZSet < Hash
+      end
+
       include Redis::Connection::CommandHelper
 
       def initialize
         @data = ExpiringHash.new
         @connected = false
@@ -121,24 +125,12 @@
       # * psubscribe
       # * publish
       # * substr
       # * unwatch
       # * watch
-      # * zadd
-      # * zcard
-      # * zcount
-      # * zincrby
-      # * zinterstore
-      # * zrange
-      # * zrangescore
-      # * zrank
-      # * zrem
       # * zremrangebyrank
       # * zremrangebyscore
-      # * zrevrange
-      # * zrevrangebyscore
-      # * zscore
       # * zunionstore
       def flushdb
         @data = ExpiringHash.new
       end
 
@@ -218,17 +210,18 @@
       end
 
       def hget(key, field)
         return unless @data[key]
         fail "Not a hash" unless @data[key].is_a?(Hash)
-        @data[key][field]
+        @data[key][field.to_s]
       end
 
       def hdel(key, field)
         return unless @data[key]
         fail "Not a hash" unless @data[key].is_a?(Hash)
         @data[key].delete(field)
+        remove_key_for_empty_collection(key)
       end
 
       def hkeys(key)
         case hash = @data[key]
           when nil then []
@@ -310,19 +303,22 @@
 
       def lrem(key, count, value)
         fail "Not a list" unless @data[key].is_a?(Array)
         return unless @data[key]
         old_size = @data[key].size
-        if count == 0
-          @data[key].delete(value)
-          old_size - @data[key].size
-        else
-          array = count > 0 ? @data[key].dup : @data[key].reverse
-          count.abs.times{ array.delete_at(array.index(value) || array.length) }
-          @data[key] = count > 0 ? array.dup : array.reverse
-          old_size - @data[key].size
-        end
+        diff =
+          if count == 0
+            @data[key].delete(value)
+            old_size - @data[key].size
+          else
+            array = count > 0 ? @data[key].dup : @data[key].reverse
+            count.abs.times{ array.delete_at(array.index(value) || array.length) }
+            @data[key] = count > 0 ? array.dup : array.reverse
+            old_size - @data[key].size
+          end
+        remove_key_for_empty_collection(key)
+        diff
       end
 
       def rpush(key, value)
         @data[key] ||= []
         fail "Not a list" unless @data[key].is_a?(Array)
@@ -382,28 +378,32 @@
       end
 
       def sadd(key, value)
         fail_unless_set(key)
         case set = @data[key]
-          when nil then @data[key] = Set.new([value.to_s])
-          when Set then set.add(value.to_s)
+          when nil then @data[key] = Set.new([value.to_s]); true
+          when Set then !!set.add?(value.to_s)
         end
       end
 
       def srem(key, value)
         fail_unless_set(key)
-        case set = @data[key]
-          when nil then return
-          when Set then set.delete(value.to_s)
-        end
+        deleted = 
+          case set = @data[key]
+            when nil then false
+            when Set then !!set.delete?(value.to_s)
+          end
+        
+        remove_key_for_empty_collection(key)
+        deleted
       end
 
       def smove(source, destination, value)
         fail_unless_set(destination)
-        if elem = self.srem(source, value)
-          self.sadd(destination, value)
-        end
+        result = self.srem(source, value)
+        self.sadd(destination, value) if result
+        result
       end
 
       def spop(key)
         fail_unless_set(key)
         elem = srandmember(key)
@@ -477,22 +477,32 @@
         end
         deleted_count = old_count - @data.keys.size
       end
 
       def setnx(key, value)
-        set(key, value) unless @data.key?(key)
+        if exists(key)
+          false
+        else
+          set(key, value)
+          true
+        end
       end
 
       def rename(key, new_key)
         return unless @data[key]
         @data[new_key] = @data[key]
         @data.expires[new_key] = @data.expires[key] if @data.expires.include?(key)
         @data.delete(key)
       end
 
       def renamenx(key, new_key)
-        rename(key, new_key) unless exists(new_key)
+        if exists(new_key)
+          false
+        else
+          rename(key, new_key)
+          true
+        end
       end
 
       def expire(key, ttl)
         return unless @data[key]
         @data.expires[key] = Time.now + ttl
@@ -511,38 +521,41 @@
         @data.expires[key] = Time.at(timestamp)
         true
       end
 
       def persist(key)
-        @data.expires.delete(key)
+        !!@data.expires.delete(key)
       end
 
       def hset(key, field, value)
+        field = field.to_s
         case hash = @data[key]
-          when nil then @data[key] = { field => value.to_s }
-          when Hash then hash[field] = value.to_s
+          when nil then @data[key] = { field => value.to_s }; true
+          when Hash then result = !hash.include?(field); hash[field] = value.to_s; result
           else fail "Not a hash"
         end
       end
 
       def hsetnx(key, field, value)
-        return if (@data[key][field] rescue false)
+        field = field.to_s
+        return false if (@data[key][field] rescue false)
         hset(key, field, value)
       end
 
       def hmset(key, *fields)
         @data[key] ||= {}
-        fail "Not a hash" unless @data[key].is_a?(Hash) 
+        fail "Not a hash" unless @data[key].is_a?(Hash)
         fields.each_slice(2) do |field|
           @data[key][field[0].to_s] = field[1].to_s
         end
       end
 
       def hmget(key, *fields)
         values = []
         fields.each do |field|
-          case hash = @data[key] 
+          field = field.to_s
+          case hash = @data[key]
             when nil then values << nil
             when Hash then values << hash[field]
             else fail "Not a hash"
           end
         end
@@ -565,14 +578,15 @@
         end
       end
 
       def hincrby(key, field, increment)
         case hash = @data[key]
-          when nil then @data[key] = { field => value.to_s }
+          when nil then @data[key] = { field => increment.to_s }
           when Hash then hash[field] = (hash[field].to_i + increment.to_i).to_s
           else fail "Not a hash"
         end
+        @data[key][field].to_i
       end
 
       def hexists(key, field)
         return unless @data[key]
         fail "Not a hash" unless @data[key].is_a?(Hash)
@@ -602,11 +616,11 @@
         old_val.each_slice(8){|b| new_val = new_val + b.join("").to_i(2).chr }
         @data[key] = new_val
       end
 
       def setex(key, seconds, value)
-        @data[key] = value
+        @data[key] = value.to_s
         expire(key, seconds)
       end
 
       def setrange(key, offset, value)
         return unless @data[key]
@@ -639,25 +653,29 @@
       end
 
       def incr(key)
         @data[key] = (@data[key] || "0")
         @data[key] = (@data[key].to_i + 1).to_s
+        @data[key].to_i
       end
 
       def incrby(key, by)
         @data[key] = (@data[key] || "0")
         @data[key] = (@data[key].to_i + by.to_i).to_s
+        @data[key].to_i
       end
 
       def decr(key)
         @data[key] = (@data[key] || "0")
         @data[key] = (@data[key].to_i - 1).to_s
+        @data[key].to_i
       end
 
       def decrby(key, by)
         @data[key] = (@data[key] || "0")
         @data[key] = (@data[key].to_i - by.to_i).to_s
+        @data[key].to_i
       end
 
       def type(key)
         case value = @data[key]
           when nil then "none"
@@ -684,16 +702,151 @@
         @buffer = []
         yield if block_given?
         "OK"
       end
 
+      def zadd(key, score, value)
+        fail_unless_zset(key)
+        @data[key] ||= ZSet.new
+        exists = @data[key].key?(value.to_s)
+        @data[key][value.to_s] = score.to_i
+        !exists
+      end
+
+      def zrem(key, value)
+        fail_unless_zset(key)
+        exists = false
+        exists = @data[key].delete(value.to_s) if @data[key]
+        remove_key_for_empty_collection(key)
+        !!exists
+      end
+
+      def zcard(key)
+        fail_unless_zset(key)
+        @data[key] ? @data[key].size : 0
+      end
+
+      def zscore(key, value)
+        fail_unless_zset(key)
+        @data[key] && @data[key][value.to_s].to_s
+      end
+
+      def zcount(key, min, max)
+        fail_unless_zset(key)
+        return 0 unless @data[key]
+        zrange_select_by_score(key, min, max).size
+      end
+
+      def zincrby(key, num, value)
+        fail_unless_zset(key)
+        @data[key][value.to_s] ||= 0
+        @data[key][value.to_s] += num
+        @data[key][value.to_s].to_s
+      end
+
+      def zrank(key, value)
+        fail_unless_zset(key)
+        @data[key].keys.sort_by {|k| @data[key][k] }.index(value.to_s)
+      end
+
+      def zrevrank(key, value)
+        fail_unless_zset(key)
+        @data[key].keys.sort_by {|k| -@data[key][k] }.index(value.to_s)
+      end
+
+      def zrange(key, start, stop, with_scores = nil)
+        fail_unless_zset(key)
+        return [] unless @data[key]
+
+        if with_scores
+          @data[key].sort_by {|_,v| v }
+        else
+          @data[key].keys.sort_by {|k| @data[key][k] }
+        end[start..stop].flatten.map(&:to_s)
+      end
+
+      def zrevrange(key, start, stop, with_scores = nil)
+        fail_unless_zset(key)
+        return [] unless @data[key]
+
+        if with_scores
+          @data[key].sort_by {|_,v| -v }
+        else
+          @data[key].keys.sort_by {|k| -@data[key][k] }
+        end[start..stop].flatten.map(&:to_s)
+      end
+
+      def zrangebyscore(key, min, max, with_scores = nil)
+        fail_unless_zset(key)
+        return [] unless @data[key]
+
+        range = zrange_select_by_score(key, min, max)
+        if with_scores
+          range.sort_by {|_,v| v }
+        else
+          range.keys.sort_by {|k| range[k] }
+        end.flatten.map(&:to_s)
+      end
+
+      def zrevrangebyscore(key, max, min, with_scores = nil)
+        fail_unless_zset(key)
+        return [] unless @data[key]
+
+        range = zrange_select_by_score(key, min, max)
+        if with_scores
+          range.sort_by {|_,v| -v }
+        else
+          range.keys.sort_by {|k| -range[k] }
+        end.flatten.map(&:to_s)
+      end
+
+      def zinterstore(out, _, *keys)
+        fail_unless_zset(out)
+
+        hashes = keys.map do |src|
+          case @data[src]
+          when Set
+            Hash[@data[src].zip([0] * @data[src].size)]
+          when Hash
+            @data[src]
+          else
+            {}
+          end
+        end
+
+        @data[out] = ZSet.new
+        values = hashes.inject([]) {|r, h| r.empty? ? h.keys : r & h.keys }
+        values.each do |value|
+          @data[out][value] = hashes.inject(0) {|n, h| n + h[value].to_i }
+        end
+
+        @data[out].size
+      end
+
       private
+
         def is_a_set?(key)
           @data[key].is_a?(Set) || @data[key].nil?
         end
 
         def fail_unless_set(key)
           fail "Not a set" unless is_a_set?(key)
+        end
+
+        def is_a_zset?(key)
+          @data[key].is_a?(ZSet) || @data[key].nil?
+        end
+
+        def fail_unless_zset(key)
+          fail "Not a sorted set" unless is_a_zset?(key)
+        end
+
+        def zrange_select_by_score(key, min, max)
+          @data[key].reject {|_,v| v < min || v > max }
+        end
+
+        def remove_key_for_empty_collection(key)
+          del(key) if @data[key] && @data[key].empty?
         end
     end
   end
 end