From 9b448040a332bb7273f73aa1005fe6922953ff62 Mon Sep 17 00:00:00 2001
From: Daniel Pepper <pepper.daniel@gmail.com>
Date: Tue, 19 Nov 2024 06:33:56 +0800
Subject: [PATCH] rack cleanup

---
 Gemfile.lock                   |  4 ++++
 lib/singed/rack_middleware.rb  |  8 ++-----
 singed.gemspec                 |  1 +
 spec/singed/middleware_spec.rb | 41 ++++++++++++++++++++++++++++++++++
 spec/spec_helper.rb            |  4 ++++
 5 files changed, 52 insertions(+), 6 deletions(-)
 create mode 100644 spec/singed/middleware_spec.rb

diff --git a/Gemfile.lock b/Gemfile.lock
index 26a6bb8..dff77e1 100644
--- a/Gemfile.lock
+++ b/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!
diff --git a/lib/singed/rack_middleware.rb b/lib/singed/rack_middleware.rb
index bce1e2f..f267974 100644
--- a/lib/singed/rack_middleware.rb
+++ b/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)
diff --git a/singed.gemspec b/singed.gemspec
index 5f577ba..2acde82 100644
--- a/singed.gemspec
+++ b/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
diff --git a/spec/singed/middleware_spec.rb b/spec/singed/middleware_spec.rb
new file mode 100644
index 0000000..8e09a03
--- /dev/null
+++ b/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_any_instance_of(Object).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_any_instance_of(Object).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
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 1850212..bdb02d9 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -96,4 +96,8 @@
   #   # test failures related to randomization by passing the same `--seed` value
   #   # as the one that triggered the failure.
   #   Kernel.srand config.seed
+
+  config.before do
+    Singed.output_directory = '/tmp/speedscope'
+  end
 end