From 19581d6e3ea76472fc64c62475b41c0924a582b4 Mon Sep 17 00:00:00 2001
From: Sdogruyol <dogruyolserdar@gmail.com>
Date: Sun, 3 Jul 2016 11:13:11 +0200
Subject: [PATCH] Improve common_exception_handler to handler custom 500
 errors. Fixes #175.

---
 spec/common_exception_handler_spec.cr | 39 ++++++++++++++++++---------
 src/kemal/common_exception_handler.cr | 19 +++++++------
 2 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/spec/common_exception_handler_spec.cr b/spec/common_exception_handler_spec.cr
index 488145f..6593959 100644
--- a/spec/common_exception_handler_spec.cr
+++ b/spec/common_exception_handler_spec.cr
@@ -1,16 +1,22 @@
 require "./spec_helper"
 
 describe "Kemal::CommonExceptionHandler" do
-  # it "renders 404 on route not found" do
-  #   get "/" do |env|
-  #     "Hello"
-  #   end
-  #
-  #   request = HTTP::Request.new("GET", "/asd")
-  #   client_response = call_request_on_app(request)
-  #   client_response.status_code.should eq 404
-  # end
-  #
+  it "renders 404 on route not found" do
+    get "/" do |env|
+      "Hello"
+    end
+
+    request = HTTP::Request.new("GET", "/asd")
+    io = MemoryIO.new
+    response = HTTP::Server::Response.new(io)
+    context = HTTP::Server::Context.new(request, response)
+    Kemal::CommonExceptionHandler::INSTANCE.call(context)
+    response.close
+    io.rewind
+    response = HTTP::Client::Response.from_io(io, decompress: false)
+    response.status_code.should eq 404
+  end
+
   # it "renders custom error" do
   #   error 403 do
   #     "403 error"
@@ -21,8 +27,15 @@ describe "Kemal::CommonExceptionHandler" do
   #   end
   #
   #   request = HTTP::Request.new("GET", "/")
-  #   client_response = call_request_on_app(request)
-  #   client_response.status_code.should eq 403
-  #   client_response.body.should eq "403 error"
+  #   io = MemoryIO.new
+  #   response = HTTP::Server::Response.new(io)
+  #   context = HTTP::Server::Context.new(request, response)
+  #   Kemal::RouteHandler::INSTANCE.call(context)
+  #   Kemal::CommonExceptionHandler::INSTANCE.call(context)
+  #   response.close
+  #   io.rewind
+  #   response = HTTP::Client::Response.from_io(io, decompress: false)
+  #   response.status_code.should eq 403
+  #   response.body.should eq "403 error"
   # end
 end
diff --git a/src/kemal/common_exception_handler.cr b/src/kemal/common_exception_handler.cr
index 8d4bccf..8893c4f 100644
--- a/src/kemal/common_exception_handler.cr
+++ b/src/kemal/common_exception_handler.cr
@@ -7,21 +7,24 @@ module Kemal
         call_next(context)
       rescue Kemal::Exceptions::RouteNotFound
         context.response.content_type = "text/html"
-        context.response.status_code = 404
-        context.response.print Kemal.config.error_handlers[404].call(context)
-        return context
+        call_exception_with_status_code(context, 404)
       rescue Kemal::Exceptions::CustomException
-        status_code = context.response.status_code
-        if Kemal.config.error_handlers.has_key?(status_code)
-          context.response.print Kemal.config.error_handlers[status_code].call(context)
-          return context
-        end
+        call_exception_with_status_code(context, context.response.status_code)
       rescue ex : Exception
         context.response.content_type = "text/html"
         Kemal.config.logger.write("Exception: #{ex.inspect_with_backtrace}\n")
+        return call_exception_with_status_code(context, 500) if Kemal.config.error_handlers.has_key?(500)
         verbosity = Kemal.config.env == "production" ? false : true
         return render_500(context, ex.inspect_with_backtrace, verbosity)
       end
     end
+
+    def call_exception_with_status_code(context, status_code)
+      if Kemal.config.error_handlers.has_key?(status_code)
+        context.response.status_code = status_code
+        context.response.print Kemal.config.error_handlers[status_code].call(context)
+        return context
+      end
+    end
   end
 end