Add support to each_batch (#2)

* Add support to each_batch

* Allow to set key on deliver_message

* require batch

lib/fake/kafka.rb

@@ -1,7 +1,8 @@
 require 'fake/kafka/version'
+require 'fake/kafka/batch'
+require 'fake/kafka/message'
 require 'fake/kafka/consumer'
 require 'fake/kafka/producer'
-require 'fake/kafka/message'
 
 module Fake
   class Kafka
@@ -17,8 +18,8 @@ def paused?(topic, partition)
       !!@paused_partitions[topic][partition]
     end
 
-    def deliver_message(value, topic:)
-      @messages << Message.new(value, nil, topic, 0, 0)
+    def deliver_message(value, topic:, key: nil)
+      @messages << Message.new(value, key, topic, 0, 0)
     end
 
     def messages_in(topic)

lib/fake/kafka/batch.rb

@@ -0,0 +1,10 @@
+class Fake::Kafka::Batch
+  attr_reader :topic, :partition, :messages, :highwater_mark_offset
+
+  def initialize(topic:, partition:, messages:, highwater_mark_offset:)
+    @topic = topic
+    @partition = partition
+    @messages = messages
+    @highwater_mark_offset = highwater_mark_offset
+  end
+end

lib/fake/kafka/consumer.rb

@@ -20,6 +20,20 @@ def each_message(*options, &block)
     end
   end
 
+  def each_batch(*options, &block)
+    begin
+      batch = Fake::Kafka::Batch.new(
+        topic: @kafka.messages.first.topic,
+        partition: @kafka.messages.first.partition,
+        messages: @kafka.messages,
+        highwater_mark_offset: @kafka.messages.first.offset
+      )
+      block.call(batch)
+    rescue StandardError => e
+      raise Kafka::ProcessingError.new(batch.topic, batch.partition, batch.highwater_mark_offset)
+    end
+  end
+
   def pause(topic, partition, timeout:, max_timeout: nil, exponential_backoff: false)
     @kafka.paused_partitions[topic] ||= {}
     @kafka.paused_partitions[topic][partition] = true