diff --git a/Gemfile.lock b/Gemfile.lock
index 7f7dc71..0cb6e40 100644
--- a/Gemfile.lock
+++ b/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
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 56936ea..c93ba4a 100644
--- a/singed.gemspec
+++ b/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
diff --git a/spec/singed/middleware_spec.rb b/spec/singed/middleware_spec.rb
new file mode 100644
index 0000000..0fa2089
--- /dev/null
+++ b/spec/singed/middleware_spec.rb
@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require "tmpdir"
+
+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) { {} }
+
+  before do
+    allow_any_instance_of(Singed::Flamegraph).to receive(:open)
+    Singed.output_directory = Dir.mktmpdir("singed-spec")
+  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 "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).and_call_original
+      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