module MockRedis::ZsetMethods

Public Instance Methods

zadd(key, *args) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 10
def zadd(key, *args)
  zadd_options = {}
  zadd_options = args.pop if args.last.is_a?(Hash)

  if zadd_options&.include?(:nx) && zadd_options&.include?(:xx)
    raise Redis::CommandError, 'ERR XX and NX options at the same time are not compatible'
  end

  if args.size == 1 && args[0].is_a?(Array)
    zadd_multiple_members(key, args.first, zadd_options)
  elsif args.size == 2
    score, member = args
    zadd_one_member(key, score, member, zadd_options)
  else
    raise Redis::CommandError, 'ERR wrong number of arguments'
  end
end
zcard(key) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 89
def zcard(key)
  with_zset_at(key, &:size)
end
zcount(key, min, max) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 93
def zcount(key, min, max)
  assert_range_args(min, max)

  with_zset_at(key) do |zset|
    zset.in_range(min, max).size
  end
end
zincrby(key, increment, member) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 101
def zincrby(key, increment, member)
  assert_scorey(increment)
  member = member.to_s
  with_zset_at(key) do |z|
    old_score = z.include?(member) ? z.score(member) : 0
    new_score = old_score + increment
    z.add(new_score, member)
    new_score.to_f
  end
end
zinterstore(destination, keys, options = {}) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 112
def zinterstore(destination, keys, options = {})
  assert_has_args(keys, 'zinterstore')

  data[destination] = combine_weighted_zsets(keys, options, :intersection)
  zcard(destination)
end
zpopmax(key, count = 1) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 167
def zpopmax(key, count = 1)
  with_zset_at(key) do |z|
    pairs = z.sorted.reverse.first(count)
    pairs.each { |pair| z.delete?(pair.last) }
    retval = to_response(pairs, with_scores: true)
    count == 1 ? retval.first : retval
  end
end
zpopmin(key, count = 1) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 158
def zpopmin(key, count = 1)
  with_zset_at(key) do |z|
    pairs = z.sorted.first(count)
    pairs.each { |pair| z.delete?(pair.last) }
    retval = to_response(pairs, with_scores: true)
    count == 1 ? retval.first : retval
  end
end
zrange(key, start, stop, options = {}) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 119
def zrange(key, start, stop, options = {})
  with_zset_at(key) do |z|
    start = [start.to_i, -z.sorted.size].max
    stop = stop.to_i
    to_response(z.sorted[start..stop] || [], options)
  end
end
zrangebyscore(key, min, max, options = {}) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 127
def zrangebyscore(key, min, max, options = {})
  assert_range_args(min, max)

  with_zset_at(key) do |zset|
    all_results = zset.in_range(min, max)
    to_response(apply_limit(all_results, options[:limit]), options)
  end
end
zrank(key, member) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 136
def zrank(key, member)
  with_zset_at(key) { |z| z.sorted_members.index(member.to_s) }
end
zrem(key, *args) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 140
def zrem(key, *args)
  if !args.first.is_a?(Array)
    retval = with_zset_at(key) { |z| !!z.delete?(args.first.to_s) }
  else
    args = args.first
    if args.empty?
      raise Redis::CommandError, "ERR wrong number of arguments for 'zrem' command"
    else
      retval = args.map { |member| !!zscore(key, member.to_s) }.count(true)
      with_zset_at(key) do |z|
        args.each { |member| z.delete?(member) }
      end
    end
  end

  retval
end
zremrangebyrank(key, start, stop) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 182
def zremrangebyrank(key, start, stop)
  zrange(key, start, stop).
    each { |member| zrem(key, member) }.
    size
end
zremrangebyscore(key, min, max) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 188
def zremrangebyscore(key, min, max)
  assert_range_args(min, max)

  zrangebyscore(key, min, max).
    each { |member| zrem(key, member) }.
    size
end
zrevrange(key, start, stop, options = {}) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 176
def zrevrange(key, start, stop, options = {})
  with_zset_at(key) do |z|
    to_response(z.sorted.reverse[start..stop] || [], options)
  end
end
zrevrangebyscore(key, max, min, options = {}) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 196
def zrevrangebyscore(key, max, min, options = {})
  assert_range_args(min, max)

  with_zset_at(key) do |zset|
    to_response(
      apply_limit(
        zset.in_range(min, max).reverse,
        options[:limit]
      ),
      options
    )
  end
end
zrevrank(key, member) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 210
def zrevrank(key, member)
  with_zset_at(key) { |z| z.sorted_members.reverse.index(member.to_s) }
end
zscan(key, cursor, opts = {}) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 214
def zscan(key, cursor, opts = {})
  opts = opts.merge(key: lambda { |x| x[0] })
  common_scan(zrange(key, 0, -1, withscores: true), cursor, opts)
end
zscan_each(key, opts = {}, &block) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 219
def zscan_each(key, opts = {}, &block)
  return to_enum(:zscan_each, key, opts) unless block_given?
  cursor = 0
  loop do
    cursor, values = zscan(key, cursor, opts)
    values.each(&block)
    break if cursor == '0'
  end
end
zscore(key, member) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 229
def zscore(key, member)
  with_zset_at(key) do |z|
    score = z.score(member.to_s)
    score&.to_f
  end
end
zunionstore(destination, keys, options = {}) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 236
def zunionstore(destination, keys, options = {})
  assert_has_args(keys, 'zunionstore')

  data[destination] = combine_weighted_zsets(keys, options, :union)
  zcard(destination)
end

Private Instance Methods

apply_limit(collection, limit) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 245
def apply_limit(collection, limit)
  if limit
    if limit.is_a?(Array) && limit.length == 2
      offset, count = limit
      collection.drop(offset).take(count)
    else
      raise Redis::CommandError, 'ERR syntax error'
    end
  else
    collection
  end
end
assert_coercible_zsety(key) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 341
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
assert_range_args(min, max) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 362
def assert_range_args(min, max)
  [min, max].each do |value|
    assert_scorey(value, 'ERR min or max is not a float')
  end
end
assert_scorey(value, message = 'ERR value is not a valid float') click to toggle source
# File lib/mock_redis/zset_methods.rb, line 353
def assert_scorey(value, message = 'ERR value is not a valid float')
  return if value.to_s =~ /\(?(\-|\+)inf/

  value = $1 if value.to_s =~ /\((.*)/
  unless looks_like_float?(value)
    raise Redis::CommandError, message
  end
end
assert_zsety(key) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 334
def assert_zsety(key)
  unless zsety?(key)
    raise Redis::CommandError,
      'WRONGTYPE Operation against a key holding the wrong kind of value'
  end
end
coerce_to_zset(set) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 296
def coerce_to_zset(set)
  zset = Zset.new
  set.each do |member|
    zset.add(1.0, member)
  end
  zset
end
coercible_zsety?(key) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 330
def coercible_zsety?(key)
  zsety?(key) || data[key].is_a?(Set)
end
combine_weighted_zsets(keys, options, how) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 268
def combine_weighted_zsets(keys, options, how)
  weights = options.fetch(:weights, keys.map { 1 })
  if weights.length != keys.length
    raise Redis::CommandError, 'ERR syntax error'
  end

  aggregator = case options.fetch(:aggregate, :sum).to_s.downcase.to_sym
               when :sum
                 proc { |a, b| [a, b].compact.reduce(&:+) }
               when :min
                 proc { |a, b| [a, b].compact.min }
               when :max
                 proc { |a, b| [a, b].compact.max }
               else
                 raise Redis::CommandError, 'ERR syntax error'
               end

  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)
      end
    end.reduce do |za, zb|
      za.send(how, zb, &aggregator)
    end
  end
end
looks_like_float?(x) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 348
def looks_like_float?(x)
  # ugh, exceptions for flow control.
  !!Float(x) rescue false
end
to_response(score_member_pairs, options) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 258
def to_response(score_member_pairs, options)
  score_member_pairs.map do |(score, member)|
    if options[:with_scores] || options[:withscores]
      [member, score.to_f]
    else
      member
    end
  end
end
with_zset_at(key, coercible: false, &blk) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 304
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
with_zsets_at(*keys, coercible: false) { |*([set] + sets)| ... } click to toggle source
# File lib/mock_redis/zset_methods.rb, line 314
def with_zsets_at(*keys, coercible: false, &blk)
  if keys.length == 1
    with_zset_at(keys.first, coercible: coercible, &blk)
  else
    with_zset_at(keys.first, coercible: coercible) do |set|
      with_zsets_at(*(keys[1..-1]), coercible: coercible) do |*sets|
        yield(*([set] + sets))
      end
    end
  end
end
zadd_multiple_members(key, args, zadd_options = {}) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 59
def zadd_multiple_members(key, args, zadd_options = {})
  assert_has_args(args, 'zadd')

  args = args.each_slice(2).to_a unless args.first.is_a?(Array)
  with_zset_at(key) do |zset|
    if zadd_options[:incr]
      raise Redis::CommandError, 'ERR INCR option supports a single increment-element pair'
    elsif zadd_options[:xx]
      args.each { |score, member| zset.include?(member) && zset.add(score, member.to_s) }
      0
    elsif zadd_options[:nx]
      args.reduce(0) do |retval, (score, member)|
        unless zset.include?(member)
          zset.add(score, member.to_s)
          retval += 1
        end
        retval
      end
    else
      args.reduce(0) do |retval, (score, member)|
        retval += 1 unless zset.include?(member)
        zset.add(score, member.to_s)
        retval
      end
    end
  end
end
zadd_one_member(key, score, member, zadd_options = {}) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 28
def zadd_one_member(key, score, member, zadd_options = {})
  assert_scorey(score) unless score.to_s =~ /(\+|\-)inf/

  with_zset_at(key) do |zset|
    if zadd_options[:incr]
      if zadd_options[:xx]
        member_present = zset.include?(member)
        return member_present ? zincrby(key, score, member) : nil
      end

      if zadd_options[:nx]
        member_present = zset.include?(member)
        return member_present ? nil : zincrby(key, score, member)
      end

      zincrby(key, score, member)
    elsif zadd_options[:xx]
      zset.add(score, member.to_s) if zset.include?(member)
      false
    elsif zadd_options[:nx]
      !zset.include?(member) && !!zset.add(score, member.to_s)
    else
      retval = !zscore(key, member)
      zset.add(score, member.to_s)
      retval
    end
  end
end
zsety?(key) click to toggle source
# File lib/mock_redis/zset_methods.rb, line 326
def zsety?(key)
  data[key].nil? || data[key].is_a?(Zset)
end