rack cleanup

Author: dpep

Gemfile.lock

@@ -48,6 +48,8 @@ GEM
     prism (1.9.0)
     racc (1.7.3)
     rack (3.2.5)
+    rack-test (2.2.0)
+      rack (>= 1.3)
     rainbow (3.1.1)
     rake (13.0.6)
     redis-client (0.26.4)
@@ -114,6 +116,7 @@ PLATFORMS
 
 DEPENDENCIES
   activejob
+  rack-test
   rake (~> 13.0)
   rspec
   sidekiq

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

@@ -23,6 +23,7 @@ Gem::Specification.new do |spec|
 
   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