Merge pull request #2 from ddukbg/fix/pr-comment

Add ZstdCompressor to S3 Plugin and Fix Tests According to Maintainer's Feedback
fluent · Oct 25, 2024 · 0ad8cdf · 0ad8cdf
2 parents 6f84c65 + dd3bafb
commit 0ad8cdf
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 38 deletions.
diff --git a/lib/fluent/plugin/s3_compressor_zstd.rb b/lib/fluent/plugin/s3_compressor_zstd.rb
diff --git a/test/test_in_s3.rb b/test/test_in_s3.rb
@@ -93,7 +93,6 @@ def test_unknown_store_as
          "text" => ["text", "txt", "text/plain"],
          "gzip" => ["gzip", "gz", "application/x-gzip"],
          "gzip_command" => ["gzip_command", "gz", "application/x-gzip"],
-         "zstd" => ["zstd", "zst", "application/x-zstd"],
          "lzo" => ["lzo", "lzo", "application/x-lzop"],
          "lzma2" => ["lzma2", "xz", "application/x-xz"])
     def test_extractor(data)

diff --git a/test/test_out_s3.rb b/test/test_out_s3.rb
@@ -109,6 +109,14 @@ def test_configure_with_mime_type_lzo
       assert(e.is_a?(Fluent::ConfigError))
     end
 
+    def test_configure_with_mime_type_zstd
+      conf = CONFIG.clone
+      conf << "\nstore_as zstd\n"
+      d = create_driver(conf)
+      assert_equal 'zst', d.instance.instance_variable_get(:@compressor).ext
+      assert_equal 'application/x-zst', d.instance.instance_variable_get(:@compressor).content_type
+    end
+
     def test_configure_with_path_style
       conf = CONFIG.clone
       conf << "\nforce_path_style true\n"
@@ -456,6 +464,33 @@ def test_write_with_custom_s3_object_key_format_containing_hex_random_placeholde
     FileUtils.rm_f(s3_local_file_path)
   end
 
+  def test_write_with_zstd
+    setup_mocks(true)
+    s3_local_file_path = "/tmp/s3-test.zst"
+
+    expected_s3path = "log/events/ts=20110102-13/events_0-#{Socket.gethostname}.zst"
+
+    setup_s3_object_mocks(s3_local_file_path: s3_local_file_path, s3path: expected_s3path)
+
+    config = CONFIG_TIME_SLICE + "\nstore_as zstd\n"
+    d = create_time_sliced_driver(config)
+
+    time = event_time("2011-01-02 13:14:15 UTC")
+    d.run(default_tag: "test") do
+      d.feed(time, { "a" => 1 })
+      d.feed(time, { "a" => 2 })
+    end
+
+    File.open(s3_local_file_path, 'rb') do |file|
+      compressed_data = file.read
+      uncompressed_data = Zstd.decompress(compressed_data)
+      expected_data = %[2011-01-02T13:14:15Z\ttest\t{"a":1}\n] +
+                      %[2011-01-02T13:14:15Z\ttest\t{"a":2}\n]
+      assert_equal expected_data, uncompressed_data
+    end
+    FileUtils.rm_f(s3_local_file_path)
+  end
+
   class MockResponse
     attr_reader :data