From 9c6305c0ba8fa7017c657704dc134c6878d60a73 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Lafr=C3=A9choux?= <jerome@jolimont.fr>
Date: Thu, 2 May 2019 23:35:54 +0200
Subject: [PATCH] Let apispec manage the cast to str of response status code

This allows supporting status codes expressed as HTTPStatus.
---
 flask_rest_api/response.py |  6 +++---
 tests/test_blueprint.py    | 39 +++++++++++++++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/flask_rest_api/response.py b/flask_rest_api/response.py
index 28e0ddce..a7f80337 100644
--- a/flask_rest_api/response.py
+++ b/flask_rest_api/response.py
@@ -23,8 +23,8 @@ def response(
 
         :param schema: :class:`Schema <marshmallow.Schema>` class or instance.
             If not None, will be used to serialize response data.
-        :param int code: HTTP status code (default: 200). Used if none is
-            returned from the view function.
+        :param int|str|HTTPStatus code: HTTP status code (default: 200).
+            Used if none is returned from the view function.
         :param str description: Description of the response (default: None).
         :param dict example: Example of response message.
         :param list examples: Examples of response message.
@@ -64,7 +64,7 @@ def decorator(func):
                 resp_doc['examples'] = examples
             if headers is not None:
                 resp_doc['headers'] = headers
-            doc = {'responses': {str(code): resp_doc}}
+            doc = {'responses': {code: resp_doc}}
             func._apidoc = deepupdate(getattr(func, '_apidoc', {}), doc)
 
             @wraps(func)
diff --git a/tests/test_blueprint.py b/tests/test_blueprint.py
index d4f2e812..09407b17 100644
--- a/tests/test_blueprint.py
+++ b/tests/test_blueprint.py
@@ -1,6 +1,7 @@
 """Test Blueprint extra features"""
 
 import json
+import http
 import pytest
 
 import marshmallow as ma
@@ -431,6 +432,7 @@ def test_blueprint_doc_merged_after_prepare_doc(self, app):
         api = Api(app)
         blp = Blueprint('test', __name__, url_prefix='/test')
 
+        # This is a dummy example. In real-life, use 'example' parameter.
         doc_example = {
             'content': {'application/json': {'example': {'test': 123}}}}
 
@@ -441,7 +443,7 @@ class ItemSchema(ma.Schema):
         class Resource(MethodView):
 
             @blp.doc(**{'requestBody': doc_example})
-            @blp.doc(**{'responses': {'200': doc_example}})
+            @blp.doc(**{'responses': {200: doc_example}})
             @blp.arguments(ItemSchema)
             @blp.response(ItemSchema)
             def get(self):
@@ -452,8 +454,39 @@ def get(self):
         get = spec['paths']['/test/']['get']
         assert get['requestBody']['content']['application/json'][
             'example'] == {'test': 123}
-        assert get['responses']['200']['content']['application/json'][
-            'example'] == {'test': 123}
+        resp = get['responses']['200']
+        assert resp['content']['application/json']['example'] == {'test': 123}
+        assert 'schema' in resp['content']['application/json']
+
+    @pytest.mark.parametrize('status_code', (200, '200', http.HTTPStatus.OK))
+    def test_blueprint_response_status_code_cast_to_string(
+            self, app, status_code):
+        api = Api(app)
+        blp = Blueprint('test', __name__, url_prefix='/test')
+
+        # This is a dummy example. In real-life, use 'description' parameter.
+        doc_desc = {'description': 'Description'}
+
+        class ItemSchema(ma.Schema):
+            test = ma.fields.Int()
+
+        @blp.route('/')
+        class Resource(MethodView):
+
+            # When documenting a response, @blp.doc MUST use the same type
+            # to express the status code as the one used in @blp.response.
+            # (Default is 200 expressed as int.)
+            @blp.doc(**{'responses': {status_code: doc_desc}})
+            @blp.arguments(ItemSchema)
+            @blp.response(ItemSchema, code=status_code)
+            def get(self):
+                pass
+
+        api.register_blueprint(blp)
+        spec = api.spec.to_dict()
+        resp = spec['paths']['/test/']['get']['responses']['200']
+        assert resp['description'] == 'Description'
+        assert 'schema' in resp['content']['application/json']
 
     def test_blueprint_doc_info_from_docstring(self, app):
         api = Api(app)