From c5201932dbacc7bfd2eabf51d446d51904ec2fcd Mon Sep 17 00:00:00 2001
From: "Thomas, Jason" <jmthomas@ball.com>
Date: Tue, 30 Jun 2020 11:29:35 -0600
Subject: [PATCH] Add xread block and count support

---
 lib/mock_redis/stream.rb         |  7 +++++--
 lib/mock_redis/stream_methods.rb |  8 +++++---
 spec/commands/xread_spec.rb      | 16 ++++++++++++++++
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/lib/mock_redis/stream.rb b/lib/mock_redis/stream.rb
index 5cce8298..d5fecfed 100644
--- 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
index bb346a74..e7d57726 100644
--- 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
index 77376b81..89947bf3 100644
--- 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