simplify Rack middleware

Gemfile.lock

@@ -19,6 +19,9 @@ GEM
       ast (~> 2.4.1)
       racc
     racc (1.7.3)
+    rack (3.1.8)
+    rack-test (2.1.0)
+      rack (>= 1.3)
     rainbow (3.1.1)
     rake (13.0.6)
     regexp_parser (2.9.0)
@@ -72,6 +75,7 @@ PLATFORMS
   ruby
 
 DEPENDENCIES
+  rack-test
   rake (~> 13.0)
   rspec
   singed!

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)

singed.gemspec

@@ -24,6 +24,7 @@ Gem::Specification.new do |spec|
   spec.add_dependency "colorize"
   spec.add_dependency "stackprof", ">= 0.2.13"
 
+  spec.add_development_dependency "rack-test"
   spec.add_development_dependency "rake", "~> 13.0"
   spec.add_development_dependency "rspec"
 end

spec/singed/middleware_spec.rb

@@ -0,0 +1,41 @@
+describe Singed::RackMiddleware do
+  subject do
+    instance.call(env)
+  end
+
+  let(:app) { ->(*) { [200, {'Content-Type' => 'text/plain'}, ['OK']] } }
+  let(:instance) { described_class.new(app) }
+  let(:env) { Rack::MockRequest.env_for('/', headers) }
+  let(:headers) { {} }
+
+  it 'returns a proper rack response' do
+    status, _ = subject
+    expect(status).to eq 200
+  end
+
+  it 'does not capture a flamegraph by default' do
+    expect(instance).not_to receive(:flamegraph)
+    subject
+  end
+
+  context 'when enabled' do
+    before { allow(instance).to receive(:capture_flamegraph?).and_return(true) }
+
+    it 'captures a flamegraph' do
+      expect(instance).to receive(:flamegraph)
+      subject
+    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
+  end
+end