From 627885b17a01ef0e2607668e0fc47ef8d8b27cec Mon Sep 17 00:00:00 2001
From: Athish Pranav D <athishanna@gmail.com>
Date: Sun, 24 Nov 2024 11:51:34 +0530
Subject: [PATCH] UTs for Out-forward and Out-file

Signed-off-by: Athish Pranav D <athishanna@gmail.com>
---
 test/plugin/test_out_forward.rb | 56 +++++++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/test/plugin/test_out_forward.rb b/test/plugin/test_out_forward.rb
index 438caf8fa2..45f7f7b1cc 100644
--- a/test/plugin/test_out_forward.rb
+++ b/test/plugin/test_out_forward.rb
@@ -342,6 +342,15 @@ def try_write(chunk)
     assert_equal :gzip, node.instance_variable_get(:@compress)
   end
 
+  test 'set_compress_is_zstd' do
+    @d = d = create_driver(config + %[compress zstd])
+    assert_equal :zstd, d.instance.compress
+    assert_equal :zstd, d.instance.buffer.compress
+
+    node = d.instance.nodes.first
+    assert_equal :zstd, node.instance_variable_get(:@compress)
+  end
+
   test 'set_compress_is_gzip_in_buffer_section' do
     mock = flexmock($log)
     mock.should_receive(:log).with("buffer is compressed.  If you also want to save the bandwidth of a network, Add `compress` configuration in <match>")
@@ -359,6 +368,23 @@ def try_write(chunk)
     assert_equal :text, node.instance_variable_get(:@compress)
   end
 
+  test 'set_compress_is_zstd_in_buffer_section' do
+    mock = flexmock($log)
+    mock.should_receive(:log).with("buffer is compressed.  If you also want to save the bandwidth of a network, Add `compress` configuration in <match>")
+
+    @d = d = create_driver(config + %[
+       <buffer>
+         type memory
+         compress zstd
+       </buffer>
+     ])
+    assert_equal :text, d.instance.compress
+    assert_equal :zstd, d.instance.buffer.compress
+
+    node = d.instance.nodes.first
+    assert_equal :text, node.instance_variable_get(:@compress)
+  end
+
   test 'phi_failure_detector disabled' do
     @d = d = create_driver(config + %[phi_failure_detector false \n phi_threshold 0])
     node = d.instance.nodes.first
@@ -549,6 +575,36 @@ def try_write(chunk)
     assert_equal ['test', time, records[1]], events[1]
   end
 
+  test 'send_comprssed_message_pack_stream_if_compress_is_zstd' do
+    target_input_driver = create_target_input_driver
+
+    @d = d = create_driver(config + %[
+      flush_interval 1s
+      compress zstd
+    ])
+
+    time = event_time('2011-01-02 13:14:15 UTC')
+
+    records = [
+      {"a" => 1},
+      {"a" => 2}
+    ]
+    target_input_driver.run(expect_records: 2) do
+      d.run(default_tag: 'test') do
+        records.each do |record|
+          d.feed(time, record)
+        end
+      end
+    end
+
+    event_streams = target_input_driver.event_streams
+    assert_true event_streams[0][1].is_a?(Fluent::CompressedMessagePackEventStream)
+
+    events = target_input_driver.events
+    assert_equal ['test', time, records[0]], events[0]
+    assert_equal ['test', time, records[1]], events[1]
+  end
+
   test 'send_to_a_node_supporting_responses' do
     target_input_driver = create_target_input_driver