simplify Rack middleware

lib/singed/rack_middleware.rb

@@ -9,15 +9,11 @@ def initialize(app)
     end
 
     def call(env)
-      status, headers, body = if capture_flamegraph?(env)
-        flamegraph do
-          @app.call(env)
-        end
+      if capture_flamegraph?(env)
+        flamegraph { @app.call(env) }
       else
         @app.call(env)
       end
-
-      [status, headers, body]
     end
 
     def capture_flamegraph?(env)

spec/singed/middleware_spec.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+describe Singed::RackMiddleware do
+  subject do
+    instance.call(env)
+  end
+
+  let(:app_response) { [200, {"content-type" => "text/plain"}, ["OK"]] }
+  let(:app) { ->(*) { app_response } }
+  let(:instance) { described_class.new(app) }
+  let(:env) { Rack::MockRequest.env_for("/", headers) }
+  let(:headers) { {} }
+
+  it "returns a proper rack response" do
+    linted_app = Rack::Lint.new(instance)
+    expect { linted_app.call(env) }.not_to raise_error
+  end
+
+  it "passes through the app response unchanged" do
+    expect(subject).to eq(app_response)
+  end
+
+  context "when enabled" do
+    before do
+      allow_any_instance_of(Singed::Flamegraph).to receive(:open)
+      allow(instance).to receive(:capture_flamegraph?).and_return(true)
+    end
+
+    it "captures a flamegraph" do
+      expect(instance).to receive(:flamegraph).and_call_original
+      subject
+    end
+
+    it "returns a proper rack response" do
+      linted_app = Rack::Lint.new(instance)
+      expect { linted_app.call(env) }.not_to raise_error
+    end
+
+    it "passes through the app response unchanged" do
+      expect(subject).to eq(app_response)
+    end
+  end
+
+  describe "#capture_flamegraph?" do
+    subject { instance.capture_flamegraph?(env) }
+
+    it { is_expected.to be false }
+
+    context "when HTTP_X_SINGED is true" do
+      let(:headers) { {"HTTP_X_SINGED" => "true"} }
+
+      it { is_expected.to be true }
+    end
+
+    context "when SINGED_MIDDLEWARE_ALWAYS_CAPTURE env var is set" do
+      around do |example|
+        original = ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"]
+        described_class.remove_instance_variable(:@always_capture) if described_class.instance_variable_defined?(:@always_capture)
+        example.run
+      ensure
+        if original.nil?
+          ENV.delete("SINGED_MIDDLEWARE_ALWAYS_CAPTURE")
+        else
+          ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] = original
+        end
+        described_class.remove_instance_variable(:@always_capture) if described_class.instance_variable_defined?(:@always_capture)
+      end
+
+      %w[true 1 yes].each do |truthy_value|
+        context "when SINGED_MIDDLEWARE_ALWAYS_CAPTURE=#{truthy_value}" do
+          before { ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] = truthy_value }
+
+          it { is_expected.to be true }
+        end
+      end
+
+      context "when SINGED_MIDDLEWARE_ALWAYS_CAPTURE=false" do
+        before { ENV["SINGED_MIDDLEWARE_ALWAYS_CAPTURE"] = "false" }
+
+        it { is_expected.to be false }
+      end
+    end
+  end
+end

spec/singed_spec.rb

@@ -1,19 +1,12 @@
 # frozen_string_literal: true
 
-require "tempfile"
-
 RSpec.describe Singed do
   around do |example|
-    original_output_directory = Singed.output_directory
-    Singed.output_directory = Dir.mktmpdir("singed-spec")
     original_enabled = Singed.enabled?
-    begin
-      example.run
-    ensure
-      Singed.output_directory = original_output_directory
-      Singed.enabled = original_enabled
-      Singed.instance_variable_set(:@current_flamegraph, nil)
-    end
+    example.run
+  ensure
+    Singed.enabled = original_enabled
+    Singed.instance_variable_set(:@current_flamegraph, nil)
   end
 
   describe ".start" do
@@ -42,7 +35,6 @@
   describe ".stop" do
     before do
       Singed.enabled = true
-      Singed.output_directory = Dir.mktmpdir("singed-spec")
     end
 
     it "returns nil when not profiling" do

spec/spec_helper.rb

@@ -15,8 +15,17 @@
 # See https://rubydoc.info/gems/rspec-core/RSpec/Core/Configuration
 
 require "singed"
+require "tmpdir"
 
 RSpec.configure do |config|
+  config.around do |example|
+    Dir.mktmpdir("singed-spec") do |dir|
+      Singed.output_directory = dir
+      example.run
+    end
+  ensure
+    Singed.output_directory = nil
+  end
   # rspec-expectations config goes here. You can use an alternate
   # assertion/expectation library such as wrong or the stdlib/minitest
   # assertions if you prefer.

spec/support/sidekiq.rb

@@ -1,5 +1,4 @@
 require "singed/sidekiq"
-require "tempfile"
 
 RSpec.configure do |config|
   config.before(:suite) do
@@ -16,14 +15,6 @@
     ActiveJob::Base.queue_adapter = :sidekiq
     ActiveJob::Base.logger = Logger.new(nil)
   end
-
-  config.around(:each, sidekiq: true) do |example|
-    orig_dir = Singed.output_directory
-    Singed.output_directory = Dir.mktmpdir("singed-sidekiq-spec")
-    example.run
-  ensure
-    Singed.output_directory = orig_dir
-  end
 end
 
 # Sidekiq doesn't invoke middlewares in inline testingmode, so we need to invoke it oursleves