Allow headers to be set when generating error resp

mikeywaites · Jan 2, 2018 · 7208914 · 7208914
1 parent e07c943
commit 7208914
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 5 deletions.
diff --git a/arrested/endpoint.py b/arrested/endpoint.py
@@ -173,10 +173,32 @@ def get_request_handler(self):
 
         return self.request_handler(self, **self.get_request_handler_params())
 
-    def return_error(self, status, payload=None):
+    def return_error(self, status, payload=None, as_json=True,
+                     headers=None, content_type='application/json'):
         """Error handler called by request handlers when an error occurs and the request
         should be aborted.
 
+        By default, the error is returned using the application/json content type and
+        The payload is serialized to JSON.  If you want to return a different
+        content type you can do so using the content_type kwarg.  You can stop the
+        payload being serialized as JSON by passing as_json=False
+
+        *A special note about METHOD NOT ALLOWED.*
+
+        When return_error is called with a 405 status no response is created.  Werkzeug's
+        exception that handles this response does not allow a response body to be
+        provided.
+
+        See the link below for more info.
+        http://werkzeug.pocoo.org/docs/0.12/exceptions/#werkzeug.exceptions.MethodNotAllowed
+
+        :param status: HTTP status code for this error
+        :param payload: Response body. When as_json=True is passed to this function
+            the payload will be serialized as a JSON string.
+        :param as_json: Serialize the payload as a JSON string
+        :param headers: a dict of headers to return with response
+        :param content_type: specify the content type to return with the response
+
         Usage::
 
             def handle_post_request(self, *args, **kwargs):
@@ -189,14 +211,18 @@ def handle_post_request(self, *args, **kwargs):
 
                 return self.return_create_response()
         """
-        resp = None
-        if payload is not None:
-            payload = json.dumps(payload)
-            resp = self.make_response(payload, status=status)
 
         if status in [405]:
             abort(status)
         else:
+
+            if as_json:
+                payload = json.dumps(payload)
+
+            resp = self.make_response(
+                payload, status=status, headers=headers, mime=content_type
+            )
+
             abort(status, response=resp)
 
     def get(self, *args, **kwargs):

diff --git a/tests/test_endpoint.py b/tests/test_endpoint.py
@@ -145,6 +145,61 @@ def test_return_error_response():
         assert resp.response.data == b'{"error": true}'
 
 
+def test_return_error_response_with_headers():
+    endpoint = Endpoint()
+    try:
+        endpoint.return_error(400, payload={'error': True}, headers={'X-Test': 'foo'})
+    except HTTPException as resp:
+        assert resp.response.headers['X-Test'] == 'foo'
+
+
+def test_return_error_response_as_json_false():
+    """Ensure that when as_json=False is passed to return_error that the method will
+    not attempt to serialize the provided payload to JSON
+    """
+    endpoint = Endpoint()
+    payload = json.dumps({'error': True})
+    try:
+        endpoint.return_error(
+            400,
+            payload=payload,
+            as_json=False,
+            headers={'X-Test': 'foo'}
+        )
+    except HTTPException as resp:
+        assert resp.response.headers['X-Test'] == 'foo'
+
+
+def test_return_error_default_content_type():
+    """ensure that the default content type returned with the response is application/json
+    """
+    endpoint = Endpoint()
+    try:
+        endpoint.return_error(
+            400,
+            payload={'error': True},
+        )
+    except HTTPException as resp:
+        assert resp.response.headers['Content-Type'] == 'application/json'
+
+
+def test_return_error_custom_content_type():
+    """ensure that the a custom content type returned with the
+    response is application/json
+    """
+    endpoint = Endpoint()
+    payload = b'<html><body>Error</body></html>'
+    try:
+        endpoint.return_error(
+            400,
+            payload=payload,
+            content_type='text/html',
+            as_json=False
+        )
+    except HTTPException as resp:
+        assert resp.response.data == payload
+
+
 def test_get_calls_handle_get_request():
     class MyEndpoint(Endpoint):