Add more test for batch request

test/dummy_server.rb

@@ -73,6 +73,14 @@ def requests
     @requests
   end
 
+  def count_per_requests
+    @requests.map{|req|JSON.parse(req.body)['Records'].size}
+  end
+
+  def size_per_requests
+    @requests.map{|req|JSON.parse(req.body)['Records'].map{|r|(r['Data'] ? Base64.decode64(r['Data']).size : 0)+(r['PartitionKey'] ? r['PartitionKey'].size : 0)}.inject(:+) || 0}
+  end
+
   def raw_records
     @accepted_records
   end
@@ -136,6 +144,9 @@ def init_server
                  @error_count += 1
                  res.status = 500
                  {}
+               elsif exceeded?(req, 500, 5*1024*1024)
+                 res.status = 400
+                 {}
                else
                  put_records_boby(req)
                end
@@ -144,6 +155,9 @@ def init_server
                  @error_count += 1
                  res.status = 500
                  {}
+               elsif exceeded?(req, 500, 4*1024*1024)
+                 res.status = 400
+                 {}
                else
                  put_record_batch_boby(req)
                end
@@ -192,6 +206,13 @@ def describe_stream_boby(req)
     }
   end
 
+  def exceeded?(req, max_count, max_size)
+    records = JSON.parse(req.body)['Records']
+    count = records.size
+    size = records.map{|r|(r['Data'] ? Base64.decode64(r['Data']).size : 0)+(r['PartitionKey'] ? r['PartitionKey'].size : 0)}.inject(:+) || 0
+    count > max_count or size > max_size
+  end
+
   def put_records_boby(req)
     body = JSON.parse(req.body)
     failed_record_count = 0

test/kinesis_helper/test_api.rb

@@ -84,12 +84,13 @@ def setup
   end
 
   data(
-    'split_by_count' => [Array.new(11, {data:'a'*1}),  [10, 1]],
-    'split_by_size'  => [Array.new(11, {data:'a'*10}), [ 6, 5]],
+    'split_by_count'            => [Array.new(11, {data:'a'*1}),  [10,1]],
+    'split_by_size'             => [Array.new(11, {data:'a'*10}), [2,2,2,2,2,1]],
+    'split_by_size_with_space'  => [Array.new(11, {data:'a'*6}),  [3,3,3,2]],
   )
   def test_batch_by_limit(data)
     records, expected = data
-    result = @object.send(:batch_by_limit, records, 10, 60)
+    result = @object.send(:batch_by_limit, records, 10, 20)
     assert_equal expected, result.map(&:size)
   end
 

test/plugin/test_out_kinesis_firehose.rb

@@ -120,6 +120,27 @@ def test_max_record_size_without_append_new_line
     assert_equal 1, d.instance.log.logs.size
   end
 
+  pad = "{\"data\":\"\"}\n".size
+  data(
+    'split_by_count'           => [Array.new(501, {data:'a'}),                                        [500,1]],
+    'split_by_size'            => [Array.new(257, {data:'a'*(16384-pad)}),                            [256,1]],
+    'split_by_size_with_space' => [Array.new(255, {data:'a'*(16384-pad)})+[{data:'a'*(16384-pad+1)}], [255,1]],
+    'no_split_by_size'         => [Array.new(255, {data:'a'*(16384-pad)})+[{data:'a'*(16384-pad)}],   [256]],
+  )
+  def test_batch_request(data)
+    records, expected = data
+    d = create_driver
+    records.each do |record|
+      d.emit(record)
+    end
+    d.run
+    assert_equal records.size, @server.records.size
+    assert_equal expected, @server.count_per_requests
+    @server.size_per_requests.each do |size|
+      assert size <= 4*1024*1024
+    end
+  end
+
   def test_record_count
     @server.enable_random_error
     d = create_driver

test/plugin/test_out_kinesis_streams.rb

@@ -97,6 +97,27 @@ def test_max_record_size
     assert_equal 1, d.instance.log.logs.size
   end
 
+  pad = 32 + '{"data":""}'.size
+  data(
+    'split_by_count'           => [Array.new(501, {data:'a'}),                                        [500,1]],
+    'split_by_size'            => [Array.new(257, {data:'a'*(20480-pad)}),                            [256,1]],
+    'split_by_size_with_space' => [Array.new(255, {data:'a'*(20480-pad)})+[{data:'a'*(20480-pad+2)}], [255,1]],
+    'no_split_by_size'         => [Array.new(255, {data:'a'*(20480-pad)})+[{data:'a'*(20480-pad)}],   [256]],
+  )
+  def test_batch_request(data)
+    records, expected = data
+    d = create_driver
+    records.each do |record|
+      d.emit(record)
+    end
+    d.run
+    assert_equal records.size, @server.records.size
+    assert_equal expected, @server.count_per_requests
+    @server.size_per_requests.each do |size|
+      assert size <= 5*1024*1024
+    end
+  end
+
   def test_record_count
     @server.enable_random_error
     d = create_driver