From 45c7a5c15c7e9b95dfd8bbb87471cc0b1d3ccede Mon Sep 17 00:00:00 2001
From: namusyaka <namusyaka@gmail.com>
Date: Wed, 13 Jan 2016 10:55:21 +0900
Subject: [PATCH] Support more kindly for after method

ref #1240
---
 lib/grape/middleware/base.rb       | 13 ++++++++++++-
 spec/grape/middleware/base_spec.rb | 14 ++++++++++++--
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/lib/grape/middleware/base.rb b/lib/grape/middleware/base.rb
index ced61811f..cb847f791 100644
--- a/lib/grape/middleware/base.rb
+++ b/lib/grape/middleware/base.rb
@@ -27,7 +27,7 @@ def call!(env)
         ensure
           after_response = after
         end
-        after_response || @app_response
+        response_valid?(after_response) ? after_response : @app_response
       end
 
       # @abstract
@@ -65,6 +65,17 @@ def mime_types
         end
         types_without_params
       end
+
+      def response_valid?(response)
+        if response.is_a?(Array)
+          status, headers, body = *response
+          (status >= 100 && status < 600) && (headers.is_a?(Hash)) && body.respond_to?(:each)
+        else
+          response.is_a?(Rack::Response)
+        end
+      rescue
+        false
+      end
     end
   end
 end
diff --git a/spec/grape/middleware/base_spec.rb b/spec/grape/middleware/base_spec.rb
index 7a3a2a456..2bdf7b55e 100644
--- a/spec/grape/middleware/base_spec.rb
+++ b/spec/grape/middleware/base_spec.rb
@@ -40,11 +40,21 @@
 
   context 'after callback' do
     before do
-      allow(subject).to receive(:after).and_return([200, {}, 'Hello from after callback'])
+      allow(subject).to receive(:after).and_return([200, {}, ['Hello from after callback']])
     end
 
     it 'overwrites application response' do
-      expect(subject.call!({}).last).to eq('Hello from after callback')
+      expect(subject.call!({}).last).to eq(['Hello from after callback'])
+    end
+  end
+
+  context 'after callback with invalid response' do
+    before do
+      allow(subject).to receive(:after).and_return('invalid response')
+    end
+
+    it 'does not overwrite application response' do
+      expect(subject.call!({}).last).to eq('Hi there.')
     end
   end