Add xread block and count support (#194)

sds · Jul 2, 2020 · 09e2fdc · 09e2fdc
1 parent 094f0d7
commit 09e2fdc
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 5 deletions.
diff --git a/lib/mock_redis/stream.rb b/lib/mock_redis/stream.rb
@@ -60,9 +60,12 @@ def range(start, finish, reversed, *opts_in)
       items
     end
 
-    def read(id)
+    def read(id, *opts_in)
+      opts = options opts_in, %w[count block]
       stream_id = MockRedis::Stream::Id.new(id)
-      members.select { |m| (stream_id < m[0]) }.map { |m| [m[0].to_s, m[1]] }
+      items = members.select { |m| (stream_id < m[0]) }.map { |m| [m[0].to_s, m[1]] }
+      return items.first(opts['count'].to_i) if opts.key?('count')
+      items
     end
 
     def each

diff --git a/lib/mock_redis/stream_methods.rb b/lib/mock_redis/stream_methods.rb
@@ -66,14 +66,16 @@ def xrevrange(key, last = '+', first = '-', count: nil)
       end
     end
 
-    # TODO: Implement count and block parameters
-    def xread(keys, ids)
+    def xread(keys, ids, count: nil, block: nil)
+      args = []
+      args += ['COUNT', count] if count
+      args += ['BLOCK', block.to_i] if block
       result = {}
       keys = keys.is_a?(Array) ? keys : [keys]
       ids = ids.is_a?(Array) ? ids : [ids]
       keys.each_with_index do |key, index|
         with_stream_at(key) do |stream|
-          data = stream.read(ids[index])
+          data = stream.read(ids[index], *args)
           result[key] = data unless data.empty?
         end
       end

diff --git a/spec/commands/xread_spec.rb b/spec/commands/xread_spec.rb
@@ -47,4 +47,20 @@
     expect(@redises.xread([@key, @key1], %w[1234567891234-2 1234567891234-1]))
       .to eq({ @key1 => [['1234567891234-2', { 'key1' => 'value2' }]] })
   end
+
+  it 'supports the block parameter' do
+    @redises.xadd(@key, { key: 'value' }, id: '1234567891234-0')
+    expect(@redises.xread(@key, '0-0', block: 1000))
+      .to eq({ @key => [['1234567891234-0', { 'key' => 'value' }]] })
+  end
+
+  it 'limits results with count' do
+    @redises.xadd(@key, { key: 'value' }, id: '1234567891234-0')
+    @redises.xadd(@key, { key: 'value' }, id: '1234567891234-1')
+    @redises.xadd(@key, { key: 'value' }, id: '1234567891234-2')
+    expect(@redises.xread(@key, '0-0', count: 1))
+      .to eq({ @key => [['1234567891234-0', { 'key' => 'value' }]] })
+    expect(@redises.xread(@key, '1234567891234-0', count: 1))
+      .to eq({ @key => [['1234567891234-1', { 'key' => 'value' }]] })
+  end
 end