rack-berater 0.0.2 → 0.3.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2a5812fba3c2523be30274bcc192d67ba7cebbb1d2b52dcc0ca403d69c39630c
-  data.tar.gz: 46152472d4af26a86e506a9150ff9c56743490707b630ecddc94beb0e63f7448
+  metadata.gz: 0b726fd0709ef54111d4e767fbe040ad3ab05fd0b0426bae20c67ef41f6fc869
+  data.tar.gz: 666a443527911b61660208f5e6ebf09b2a7adf298c67cb157665988c32868623
 SHA512:
-  metadata.gz: '079407a6fd8d63c849809c972ef15739de09bd490c1e3a4046ae047dbba8b36cb7e5f0606c17b5e5a6a74ce8ec9fa95f61139afc3937a5c30643fe85ff59e91d'
-  data.tar.gz: 7d6ede65719ed13181c25aca35a48c941ee9c4777caa4f56eede0552b4acdc39d2989069659217327a0b3bc494d0c178b1ee1faa21e42f5c70dd4067f4fb077e
+  metadata.gz: f5bfe591019a79137181785ae76f6e5603b3495cc5ef8e207eab1be88bec114988bbe190509752e093aa1458ecd2ce11dd91dcd4a31fe96b9c00e94f62a931ea
+  data.tar.gz: a83f26c06eeeb1c570233b1b84bdec1dadf501e8774a388290dd6c733fd2c095435f86974b4f27216ffbd118e261206097ac3f16187dfeb2f2da8eb03adb9657

data/lib/rack/berater/prioritizer.rb

@@ -0,0 +1,93 @@
+require 'rack'
+
+module Rack
+  class Berater
+    class Prioritizer
+      ENV_KEY = 'berater_priority'
+      HEADER = 'X-Berater-Priority'
+
+      def initialize(app, options = {})
+        @app = app
+        @header = options[:header] || HEADER
+      end
+
+      def call(env)
+        priority = env[@header] || env["HTTP_#{@header.upcase.tr('-', '_')}"]
+
+        if priority
+          set_priority(priority)
+          return @app.call(env)
+        end
+
+        cache_key = cache_key_for(env)
+        cached_priority = cache_get(cache_key)
+
+        if cached_priority
+          set_priority(cached_priority)
+        end
+
+        @app.call(env).tap do |status, headers, body|
+          app_priority = headers.delete(@header) if headers
+
+          if app_priority && app_priority != cached_priority
+            # update cache for next time
+            cache_set(cache_key, app_priority)
+          end
+        end
+      ensure
+        Thread.current[ENV_KEY] = nil
+      end
+
+      def self.current_priority
+        Thread.current[ENV_KEY]
+      end
+
+      protected
+
+      def set_priority(priority)
+        Thread.current[ENV_KEY] = priority
+      end
+
+      @@cache_prefix = 'Berater:Rack:Prioritizer'
+      def cache_key_for(env)
+        req = Rack::Request.new(env)
+
+        req_method = env[Rack::REQUEST_METHOD].downcase
+        path = ''
+
+        if defined?(Rails) && Rails.respond_to?(:application) && Rails.application
+          res = Rails.application.routes.recognize_path(
+            env[Rack::PATH_INFO],
+            method: env[Rack::REQUEST_METHOD],
+          )
+          path = res.values_at(:controller, :action).compact.join('#')
+        end
+
+        if path.empty?
+          path = env['PATH_INFO'].gsub(%r{/[0-9]+(/|$)}, '/x\1')
+        end
+
+        [
+          @@cache_prefix,
+          req_method,
+          path,
+        ].join(':')
+      end
+
+      @@cache = {}
+
+      def cache_get(key)
+        synchronize { @@cache[key] }
+      end
+
+      def cache_set(key, priority)
+        synchronize { @@cache[key] = priority }
+      end
+
+      @@lock = Thread::Mutex.new
+      def synchronize(&block)
+        @@lock.synchronize(&block)
+      end
+    end
+  end
+end

data/lib/rack/berater/railtie.rb

@@ -1,9 +1,13 @@
 require 'rails/railtie'
 
 module Rack
-  module Berater
+  class Berater
     class Railtie < Rails::Railtie
-      initializer "rack.berater.initializer" do |app|
+      initializer 'rack.berater' do |app|
+        if ::Berater.middleware.include?(::Berater::Middleware::LoadShedder)
+          app.middleware.use Rack::Berater::Prioritizer
+        end
+
         app.middleware.use Rack::Berater
       end
     end

data/lib/rack/berater/version.rb

@@ -1,5 +1,5 @@
 module Rack
-  module Berater
-    VERSION = "0.0.2"
+  class Berater
+    VERSION = "0.3.1"
   end
 end

data/lib/rack/berater.rb

@@ -1,9 +1,67 @@
-require "rack/berater/version"
+require 'berater'
+require 'rack'
+require 'rack/berater/version'
+require 'set'
 
 module Rack
-  module Berater
-    autoload :Handler, "rack/berater/handler"
-    # autoload :Limiter, "rack/berater/limiter"
-    autoload :Railtie, "rack/berater/railtie"
+  class Berater
+    autoload :Prioritizer, 'rack/berater/prioritizer'
+    autoload :Railtie, 'rack/berater/railtie'
+
+    ERRORS = Set[ ::Berater::Overloaded ]
+
+    def initialize(app, options = {})
+      @app = app
+      @enabled = options[:enabled?]
+      @limiter = options[:limiter]
+      @options = {
+        headers: {},
+        status_code: options.fetch(:status_code, 429),
+      }
+
+      # configure body
+      @options[:body] = case options[:body]
+        when true, nil
+          Rack::Utils::HTTP_STATUS_CODES[@options[:status_code]]
+        when false
+          nil
+        when String
+          options[:body]
+        else
+          raise ArgumentError, 'invalid :body option: #{options[:body]}'
+        end
+
+      # configure headers
+      if @options[:body]
+        @options[:headers][Rack::CONTENT_TYPE] = 'text/plain'
+      end
+      @options[:headers].update(options.fetch(:headers, {}))
+    end
+
+    def call(env)
+      if enabled?(env)
+        limit(env) { @app.call(env) }
+      else
+        @app.call(env)
+      end
+    rescue *ERRORS => e
+      [
+        @options[:status_code],
+        @options[:headers],
+        [ @options[:body] ].compact,
+      ]
+    end
+
+    private
+
+    def enabled?(env)
+      return false unless @limiter
+      @enabled.nil? ? true : @enabled.call(env)
+    end
+
+    def limit(env, &block)
+      limiter = @limiter.respond_to?(:call) ? @limiter.call(env) : @limiter
+      limiter.limit(&block)
+    end
   end
 end

data/lib/rack-berater.rb

@@ -1,3 +1 @@
-require "berater"
-require "rack"
-require "rack/berater"
+require 'rack/berater'

data/spec/limiter_spec.rb

@@ -0,0 +1,118 @@
+describe Rack::Berater do
+  before do
+    app.use described_class, limiter: limiter, enabled?: enabled?
+  end
+  let(:limiter) { nil }
+  let(:enabled?) { nil }
+
+  let(:app) do
+    Rack::Builder.new do
+      use Rack::Lint
+      run (lambda do |env|
+        [200, {'Content-Type' => 'text/plain'}, ['OK']]
+      end)
+    end
+  end
+  let(:response) { get '/' }
+
+  shared_examples 'works nominally' do
+    it 'has the correct status code' do
+      expect(response.status).to eq 200
+    end
+
+    it 'has the correct headers' do
+      expect(response.headers).to eq({
+        'Content-Type' => 'text/plain',
+        'Content-Length' => '2',
+      })
+    end
+
+    it 'has the correct body' do
+      expect(response.body).to eq 'OK'
+    end
+  end
+
+  context 'without a limiter' do
+    before { Berater.test_mode = :fail }
+
+    include_examples 'works nominally'
+  end
+
+  describe 'limiter option' do
+    context 'when limiter is a limiter' do
+      let(:limiter) { ::Berater::Unlimiter.new }
+
+      include_examples 'works nominally'
+
+      it 'calls the limiter' do
+        expect(limiter).to receive(:limit).and_call_original
+        response
+      end
+
+      context 'when operating beyond limits' do
+        before { Berater.test_mode = :fail }
+
+        it 'returns an error' do
+          expect(response.status).to eq 429
+        end
+      end
+    end
+
+    context 'when limiter is a proc' do
+      let(:limiter_instance) { ::Berater::Unlimiter.new }
+      let(:limiter) { Proc.new { limiter_instance } }
+
+      include_examples 'works nominally'
+
+      it 'calls the proc with env' do
+        expect(limiter).to receive(:call).with(Hash).and_call_original
+        response
+      end
+
+      context 'when operating beyond limits' do
+        before { Berater.test_mode = :fail }
+
+        it 'returns an error' do
+          expect(response.status).to eq 429
+        end
+      end
+    end
+  end
+
+  describe 'enabled? option' do
+    after { expect(response.status).to eq 200 }
+
+    let(:enabled?) { double }
+
+    context 'when there is a limiter' do
+      let(:limiter) { ::Berater::Unlimiter.new }
+
+      it 'should be called with the env hash' do
+        expect(enabled?).to receive(:call) do |env|
+          expect(env).to be_a Hash
+          expect(Rack::Request.new(env).path).to eq '/'
+        end
+      end
+
+      context 'when enabled' do
+        it 'should call the limiter' do
+          expect(enabled?).to receive(:call).and_return(true)
+          expect(limiter).to receive(:limit).and_call_original
+        end
+      end
+
+      context 'when disabled' do
+        it 'should not call the limiter' do
+          expect(enabled?).to receive(:call).and_return(false)
+          expect(limiter).not_to receive(:limit)
+        end
+      end
+    end
+
+    context 'when there is no limiter' do
+      it 'should not call enabled?' do
+        expect(enabled?).not_to receive(:call)
+      end
+    end
+  end
+end

data/spec/prioritizer_spec.rb

@@ -0,0 +1,295 @@
+describe Rack::Berater::Prioritizer do
+  after do
+    cache.clear
+    Thread.current[described_class::ENV_KEY] = nil
+  end
+
+  let(:cache) { described_class.class_variable_get(:@@cache) }
+
+  describe '#call' do
+    subject { described_class.new(app) }
+
+    let(:app) { ->(*) { [200, {'Content-Type' => 'text/plain'}, ['OK']] } }
+    let(:cache_key) { subject.send(:cache_key_for, env) }
+    let(:env) { Rack::MockRequest.env_for('/') }
+
+    specify 'sanity check' do
+      expect(app).to receive(:call).and_call_original
+
+      Rack::Lint.new(subject).call(env)
+    end
+
+    it 'checks the cache' do
+      is_expected.to receive(:cache_get)
+
+      subject.call(env)
+    end
+
+    context 'with a cached priority' do
+      before do
+        allow(subject).to receive(:cache_get).with(cache_key).and_return(priority)
+      end
+
+      let(:priority) { '3' }
+
+      after { subject.call(env) }
+
+      it 'sets the priority accordingly' do
+        is_expected.to receive(:set_priority).with(priority)
+      end
+
+      it 'updates the global priority during the request' do
+        expect(app).to receive(:call) do
+          expect(described_class.current_priority).to eq priority
+        end
+      end
+
+      it 'resets the priority after the request completes' do
+        subject.call(env)
+        expect(described_class.current_priority).to be nil
+      end
+    end
+
+    context 'with an incoming priority header' do
+      let(:env) do
+        Rack::MockRequest.env_for(
+          '/',
+          described_class::HEADER => priority,
+        )
+      end
+      let(:priority) { '2' }
+
+      after { subject.call(env) }
+
+      it 'uses the header' do
+        is_expected.to receive(:set_priority).with(priority)
+      end
+
+      it 'ignores any cached value' do
+        allow(subject).to receive(:cache_get).with(cache_key).and_return('123')
+        is_expected.to receive(:set_priority).with(priority)
+      end
+
+      it 'resets the priority after the request completes' do
+        subject.call(env)
+        expect(described_class.current_priority).to be nil
+      end
+    end
+
+    context 'when the app returns a priority header' do
+      let(:app) do
+        ->(*) { [200, { described_class::HEADER => priority }, ['OK']] }
+      end
+
+      let(:priority) { '5' }
+
+      after { subject.call(env) }
+
+      it 'caches the priority' do
+        is_expected.to receive(:cache_set).with(cache_key, priority)
+      end
+
+      it 'removes the header' do
+        _, headers, _ = subject.call(env)
+        expect(headers).not_to include described_class::HEADER
+      end
+
+      it 'updates the cache when a different priority is returned' do
+        expect(subject).to receive(:cache_get).and_return('123')
+        is_expected.to receive(:cache_set).with(cache_key, priority)
+      end
+
+      it 'does not update the cache when the same priority is returned' do
+        expect(subject).to receive(:cache_get).and_return(priority)
+        is_expected.not_to receive(:cache_set)
+      end
+    end
+
+    it 'does not update the cache when no priority is returned' do
+      is_expected.not_to receive(:cache_set)
+      subject.call(env)
+    end
+  end
+
+  describe '#cache_key_for' do
+    subject{ described_class.new(nil).send(:cache_key_for, env) }
+
+    context 'with a basic env' do
+      let(:env) { Rack::MockRequest.env_for('/') }
+
+      it 'has a Berater prefix' do
+        is_expected.to match /^Berater:/
+      end
+
+      it 'combines the verb and path' do
+        is_expected.to match %r{:get:/$}
+      end
+    end
+
+    context 'with a different verb' do
+      let(:env) { Rack::MockRequest.env_for('/', method: 'PUT') }
+
+      it 'combines the verb and path' do
+        is_expected.to match %r{:put:/$}
+      end
+    end
+
+    context 'with a RESTful path' do
+      let(:env) { Rack::MockRequest.env_for('/user/123') }
+
+      it 'normalizes the id' do
+        is_expected.to match %r{:/user/x$}
+      end
+    end
+
+    context 'with a RESTful path and trailing slash' do
+      let(:env) { Rack::MockRequest.env_for('/user/123/') }
+
+      it 'normalizes the id and keeps the trailing slash' do
+        is_expected.to match %r{:/user/x/$}
+      end
+    end
+
+    context 'with a very RESTful path' do
+      let(:env) { Rack::MockRequest.env_for('/user/123/friend/456') }
+
+      it 'normalizes both ids' do
+        is_expected.to match %r{:/user/x/friend/x$}
+      end
+    end
+  end
+
+  context 'as Rack middleware' do
+    def call(path = '/')
+      get(path).body
+    end
+
+    let(:app) do
+      headers = {
+        'Content-Type' => 'text/plain',
+        described_class::HEADER => app_priority,
+      }.compact
+
+      Rack::Builder.new do
+        use Rack::Lint
+        use Rack::Berater::Prioritizer
+
+        run (lambda do |env|
+          [200, headers, [ Rack::Berater::Prioritizer.current_priority.to_s ]]
+        end)
+      end
+    end
+
+    let(:app_priority) { nil }
+
+    it 'starts empty' do
+      expect(call).to be_empty
+    end
+
+    it 'parses incoming priority header' do
+      header described_class::HEADER, '7'
+
+      expect(call).to eq '7'
+    end
+
+    context 'when app returns a priority header' do
+      let(:app_priority) { '8' }
+
+      it 'parses the priority returned from the app' do
+        expect(call).to be_empty
+        expect(cache.values).to include app_priority
+      end
+
+      it 'uses the cached priority for subsequent calls' do
+        expect(call).to be_empty
+        expect(call).to eq app_priority
+      end
+    end
+
+    # context 'when two different endpoints are called' do
+    #   fit 'parses and caches each priority' do
+    #     @app_priority = '6'
+    #     expect(call('/six')).to be_empty
+
+    #     expect(call('/six')).to eq '6'
+
+    #     @app_priority = '9'
+    #     expect(call('/nine')).to be_empty
+    #     expect(call('/nine')).to '9'
+    #   end
+    # end
+  end
+
+  context 'as Rails middleware' do
+    before do
+      class EchoController < ActionController::Base
+        def index
+          render plain: Rack::Berater::Prioritizer.current_priority
+        end
+
+        def six
+          response.set_header(Rack::Berater::Prioritizer::HEADER, '6')
+          index
+        end
+
+        def nine
+          response.set_header(Rack::Berater::Prioritizer::HEADER, '9')
+          index
+        end
+      end
+
+      Rails.application = Class.new(Rails::Application) do
+        config.eager_load = false
+        config.hosts.clear # disable hostname filtering
+      end
+      Rails.application.middleware.use described_class
+      Rails.initialize!
+
+      Rails.application.routes.draw do
+        get '/' => 'echo#index'
+        get '/six' => 'echo#six'
+        post '/nine' => 'echo#nine'
+      end
+    end
+
+    let(:app) { Rails.application }
+    let(:middleware) { described_class.new(app) }
+
+    after { Rails.application = nil }
+
+    describe '#cache_key_for' do
+      subject { described_class.new(app).send(:cache_key_for, env) }
+
+      let(:env) { Rack::MockRequest.env_for('/') }
+
+      it 'uses the controller and action name' do
+        is_expected.to match %r{:echo#index$}
+      end
+    end
+
+    context 'when a priority header is sent' do
+      before { header described_class::HEADER, priority }
+
+      let(:priority) { '6' }
+
+      it 'sets the priority' do
+        expect(get('/six').body).to eq priority
+      end
+    end
+
+    context 'when the app returns a priority' do
+      it 'does not know the first time the controller is called' do
+        expect(get('/six').body).to be_empty
+        expect(post('/nine').body).to be_empty
+      end
+
+      it 'caches the repsonses for the second time' do
+        expect(get('/six').body).to be_empty
+        expect(post('/nine').body).to be_empty
+
+        expect(get('/six').body).to eq '6'
+        expect(post('/nine').body).to eq '9'
+      end
+    end
+  end
+end

data/spec/railtie_spec.rb

@@ -0,0 +1,16 @@
+require 'rails'
+require 'rack/berater/railtie'
+
+RSpec.describe Rack::Berater::Railtie do
+  subject { Rails.initialize! }
+
+  before do
+    Rails.application = Class.new(Rails::Application) do
+      config.eager_load = false
+    end
+  end
+
+  it 'adds middleware automatically' do
+    expect(subject.middleware).to include(Rack::Berater)
+  end
+end

data/spec/rescuer_spec.rb

@@ -0,0 +1,155 @@
+describe Rack::Berater do
+  let(:app) do
+    Rack::Builder.new do
+      use Rack::Lint
+      run (lambda do |env|
+        Berater::Unlimiter() do
+          [200, {'Content-Type' => 'text/plain'}, ['OK']]
+        end
+      end)
+    end
+  end
+  let(:response) { get '/' }
+
+  shared_examples 'works nominally' do
+    it 'has the correct status code' do
+      expect(response.status).to eq 200
+    end
+
+    it 'has the correct headers' do
+      expect(response.headers).to eq({
+        'Content-Type' => 'text/plain',
+        'Content-Length' => '2',
+      })
+    end
+
+    it 'has the correct body' do
+      expect(response.body).to eq 'OK'
+    end
+  end
+
+  context 'without middleware' do
+    include_examples 'works nominally'
+
+    it 'does not catch limit errors' do
+      Berater.test_mode = :fail
+      expect {
+        response
+      }.to be_overloaded
+    end
+  end
+
+  context 'with middleware using default settings' do
+    before { app.use described_class }
+
+    include_examples 'works nominally'
+
+    it 'catches and transforms limit errors' do
+      Berater.test_mode = :fail
+      expect(response.status).to eq 429
+      expect(response.body).to eq 'Too Many Requests'
+    end
+  end
+
+  context 'with middleware using custom settings' do
+    before do
+      app.use described_class, options
+      Berater.test_mode = :fail
+    end
+
+    context 'with a custom body' do
+      context 'with body nil' do
+        let(:options) { { body: nil } }
+
+        it 'falls back to the default' do
+          expect(response.body).to eq 'Too Many Requests'
+        end
+      end
+
+      context 'with body disabled' do
+        let(:options) { { body: false } }
+
+        it 'should not send a body' do
+          expect(response.body).to be_empty
+        end
+
+        it 'should not send the Content-Type header' do
+          expect(response.headers.keys).not_to include(Rack::CONTENT_TYPE)
+        end
+      end
+
+      context 'with a string' do
+        let(:body) { 'none shall pass!' }
+        let(:options) { { body: body } }
+
+        it 'should send the custom string' do
+          expect(response.body).to eq body
+        end
+      end
+
+      context 'with an erroneous value' do
+        let(:options) { { body: 123 } }
+
+        it 'should raise an error' do
+          expect {
+            response
+          }.to raise_error(ArgumentError)
+        end
+      end
+    end
+
+    context 'with custom headers' do
+      context 'with an extra header' do
+        let(:options) { { headers: { Rack::CACHE_CONTROL => 'no-cache' } } }
+
+        it 'should contain the default headers' do
+          expect(response.headers.keys).to include(Rack::CONTENT_TYPE)
+        end
+
+        it 'should also contain the custom header' do
+          expect(response.headers).to include(options[:headers])
+        end
+      end
+
+      context 'with a new content type' do
+        let(:options) { { headers: { Rack::CONTENT_TYPE => 'application/json' } } }
+
+        it 'should override the Content-Type header' do
+          expect(response.headers).to include(options[:headers])
+        end
+      end
+    end
+
+    context 'with custom status code' do
+      let(:options) { { status_code: 503 } }
+
+      it 'catches and transforms limit errors' do
+        expect(response.status).to eq 503
+        expect(response.body).to eq 'Service Unavailable'
+      end
+    end
+  end
+
+  context 'with custom error type' do
+    before do
+      app.use described_class
+      expect(Berater::Limiter).to receive(:new).and_raise(IOError)
+    end
+
+    it 'normally crashes the app' do
+      expect { response }.to raise_error(IOError)
+    end
+
+    context 'when an error type is registered with middleware' do
+      around do |example|
+        Rack::Berater::ERRORS << IOError
+        example.run
+        Rack::Berater::ERRORS.delete(IOError)
+      end
+
+      it 'catches and transforms limit errors' do
+        expect(response.status).to eq 429
+      end
+    end
+  end
+end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: rack-berater
 version: !ruby/object:Gem::Version
-  version: 0.0.2
+  version: 0.3.1
 platform: ruby
 authors:
 - Daniel Pepper
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-09-13 00:00:00.000000000 Z
+date: 2021-11-18 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: berater
@@ -80,6 +80,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: rails
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 - !ruby/object:Gem::Dependency
   name: rspec
   requirement: !ruby/object:Gem::Requirement
@@ -116,10 +130,13 @@ extra_rdoc_files: []
 files:
 - lib/rack-berater.rb
 - lib/rack/berater.rb
-- lib/rack/berater/handler.rb
+- lib/rack/berater/prioritizer.rb
 - lib/rack/berater/railtie.rb
 - lib/rack/berater/version.rb
-- spec/handler_spec.rb
+- spec/limiter_spec.rb
+- spec/prioritizer_spec.rb
+- spec/railtie_spec.rb
+- spec/rescuer_spec.rb
 homepage: https://github.com/dpep/rack-berater
 licenses:
 - MIT
@@ -139,9 +156,12 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.0.8
+rubygems_version: 3.1.6
 signing_key:
 specification_version: 4
 summary: Rack::Berater
 test_files:
-- spec/handler_spec.rb
+- spec/railtie_spec.rb
+- spec/rescuer_spec.rb
+- spec/prioritizer_spec.rb
+- spec/limiter_spec.rb

data/lib/rack/berater/handler.rb

@@ -1,43 +0,0 @@
-module Rack
-  module Berater
-    class Handler
-      def initialize(app, options = {})
-        @app = app
-        @options = {
-          status_code: options.fetch(:status_code, 429),
-          headers: {
-            Rack::CONTENT_TYPE => "text/plain",
-          }.update(options.fetch(:headers, {})),
-          body: options.fetch(:body, true),
-        }
-      end
-
-      def call(env)
-        @app.call(env)
-      rescue ::Berater::Overloaded => e
-        code = @options[:status_code]
-
-        body = case @options[:body]
-          when true
-            Rack::Utils::HTTP_STATUS_CODES[code]
-          when nil, false
-            nil
-          when String
-            @options[:body]
-          when Proc
-            @options[:body].call(env, e)
-          else
-            raise ArgumentError, "invalid :body option: #{@options[:body]}"
-          end
-
-        headers = body ? @options[:headers] : {}
-
-        [
-          code,
-          headers,
-          [ body ].compact,
-        ]
-      end
-    end
-  end
-end

data/spec/handler_spec.rb

@@ -1,126 +0,0 @@
-describe Rack::Berater::Handler do
-  let(:app) do
-    Rack::Builder.new do
-      use Rack::Lint
-      run (lambda do |env|
-        raise ::Berater::Overloaded if ::Berater.test_mode == :fail
-        [200, {"Content-Type" => "text/plain"}, ["OK"]]
-      end)
-    end
-  end
-  let(:response) { get "/" }
-
-  shared_examples "works nominally" do
-    it "has the correct status code" do
-      expect(response.status).to eq 200
-    end
-
-    it "has the correct headers" do
-      expect(response.headers).to eq({
-        "Content-Type" => "text/plain",
-        "Content-Length" => "2",
-      })
-    end
-
-    it "has the correct body" do
-      expect(response.body).to eq "OK"
-    end
-  end
-
-  context "without Handler" do
-    include_examples "works nominally"
-
-    it "does not catch limit errors" do
-      Berater.test_mode = :fail
-      expect {
-        response
-      }.to be_overloaded
-    end
-  end
-
-  context "with Handler using default settings" do
-    context "with default settings" do
-      before { app.use described_class }
-
-      include_examples "works nominally"
-
-      it "catches and transforms limit errors" do
-        Berater.test_mode = :fail
-        expect(response.status).to eq 429
-        expect(response.body).to eq "Too Many Requests"
-      end
-    end
-  end
-
-  context "with Handler using custom settings" do
-    before do
-      app.use described_class, options
-      Berater.test_mode = :fail
-    end
-
-    context "with custom status code" do
-      let(:options) { { status_code: 503 } }
-
-      it "catches and transforms limit errors" do
-        expect(response.status).to eq 503
-        expect(response.body).to eq "Service Unavailable"
-      end
-    end
-
-    context "with body disabled" do
-      let(:options) { { body: false } }
-
-      it "should not send a body" do
-        expect(response.body).to be_empty
-      end
-
-      it "should not send the Content-Type header" do
-        expect(response.headers.keys).not_to include(Rack::CONTENT_TYPE)
-      end
-    end
-
-    context "with body nil" do
-      let(:options) { { body: nil } }
-
-      it "should not send a body" do
-        expect(response.body).to be_empty
-      end
-    end
-
-    context "with custom body" do
-      let(:body) { "none shall pass!" }
-      let(:options) { { body: body } }
-
-      it "should send the custom string" do
-        expect(response.body).to eq body
-      end
-    end
-
-    context "with a dynamic body" do
-      let(:body) { "none shall pass!" }
-      let(:fn) { proc { body } }
-      let(:options) { { body: fn } }
-
-      it "should call the Proc and send the result" do
-        expect(response.body).to eq body
-      end
-
-      it "should pass in the env and error" do
-        expect(fn).to receive(:call).with(Hash, ::Berater::Overloaded)
-        response
-      end
-    end
-
-    context "with custom headers" do
-      let(:options) { { headers: { Rack::CACHE_CONTROL => "no-cache" } } }
-
-      it "should contain the default headers" do
-        expect(response.headers.keys).to include(Rack::CONTENT_TYPE)
-      end
-
-      it "should also contain custom header" do
-        expect(response.headers).to include(options[:headers])
-      end
-    end
-  end
-end