diff --git a/test/dummy_server.rb b/test/dummy_server.rb index b4509b8..6ace8e5 100644 --- a/test/dummy_server.rb +++ b/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 diff --git a/test/kinesis_helper/test_api.rb b/test/kinesis_helper/test_api.rb index 086a34e..cebf53e 100644 --- a/test/kinesis_helper/test_api.rb +++ b/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 diff --git a/test/plugin/test_out_kinesis_firehose.rb b/test/plugin/test_out_kinesis_firehose.rb index 465d7a9..354a0d1 100644 --- a/test/plugin/test_out_kinesis_firehose.rb +++ b/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 diff --git a/test/plugin/test_out_kinesis_streams.rb b/test/plugin/test_out_kinesis_streams.rb index 2199596..f44037c 100644 --- a/test/plugin/test_out_kinesis_streams.rb +++ b/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