Merge pull request #164 from jhawthorn/support-writing-result-to-stringio

Author: joshuay03

lib/vernier/result.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Vernier
   class Result
     attr_accessor :stack_table
@@ -41,10 +43,17 @@ def to_cpuprofile
     def write(out:, format: "firefox")
       case format
       when "cpuprofile"
-        File.binwrite(out, to_cpuprofile)
-      when nil, "firefox"
-        gzip = out.end_with?(".gz")
-        File.binwrite(out, to_firefox(gzip:))
+        if out.respond_to?(:write)
+          out.write(to_cpuprofile)
+        else
+          File.binwrite(out, to_cpuprofile)
+        end
+      when "firefox", nil
+        if out.respond_to?(:write)
+          out.write(to_firefox)
+        else
+          File.binwrite(out, to_firefox(gzip: out.end_with?(".gz")))
+        end
       else
         raise ArgumentError, "unknown format: #{format}"
       end

test/test_result.rb

@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require "tempfile"
+require "stringio"
+require "zlib"
+
+class TestResult < Minitest::Test
+  include FirefoxTestHelpers
+  include CpuprofileTestHelpers
+
+  def setup
+    @result = Vernier.trace do
+      sleep 0.001
+    end
+  end
+
+  def test_write_to_file_firefox_format
+    Tempfile.open("vernier_test") do |tempfile|
+      @result.write(out: tempfile.path, format: "firefox")
+      tempfile.rewind
+      content = tempfile.read
+      assert_valid_firefox_profile(content)
+    end
+  end
+
+  def test_write_to_file_cpuprofile_format
+    Tempfile.open("vernier_test") do |tempfile|
+      @result.write(out: tempfile.path, format: "cpuprofile")
+      tempfile.rewind
+      content = tempfile.read
+      assert_valid_cpuprofile(content)
+    end
+  end
+
+  def test_write_to_stringio_firefox_format
+    stringio = StringIO.new
+    @result.write(out: stringio, format: "firefox")
+    stringio.rewind
+    content = stringio.read
+    assert_valid_firefox_profile(content)
+  end
+
+  def test_write_to_stringio_cpuprofile_format
+    stringio = StringIO.new
+    @result.write(out: stringio, format: "cpuprofile")
+    stringio.rewind
+    content = stringio.read
+    assert_valid_cpuprofile(content)
+  end
+
+  def test_write_to_file_with_gz_extension
+    Tempfile.open(["vernier_test", ".gz"]) do |tempfile|
+      @result.write(out: tempfile.path, format: "firefox")
+      tempfile.rewind
+      content = tempfile.read
+      decompressed = Zlib.gunzip(content)
+      assert_valid_firefox_profile(decompressed)
+    end
+  end
+
+  def test_write_to_stringio_ignores_gz_extension_logic
+    stringio = StringIO.new
+    @result.write(out: stringio, format: "firefox")
+    stringio.rewind
+    content = stringio.read
+    assert_valid_firefox_profile(content)
+  end
+
+  def test_write_with_invalid_format_raises_error
+    stringio = StringIO.new
+    error = assert_raises(ArgumentError) do
+      @result.write(out: stringio, format: "invalid")
+    end
+    assert_equal "unknown format: invalid", error.message
+  end
+end