diff --git a/tests/test_flask.py b/tests/test_flask.py index da9e1c6..e5add1f 100644 --- a/tests/test_flask.py +++ b/tests/test_flask.py @@ -1,6 +1,7 @@ import unittest from flask import Flask, json from wsmeext.flask import signature +from wsme.api import Response from wsme.types import Base, text @@ -51,6 +52,18 @@ def post_model(body): return Model(name=body.name) +@test_app.route('/status_sig') +@signature(int, status_code=201) +def get_status_sig(): + return 1 + + +@test_app.route('/status_response') +@signature(int) +def get_status_response(): + return Response(1, status_code=201) + + class FlaskrTestCase(unittest.TestCase): def setUp(self): @@ -90,6 +103,14 @@ class FlaskrTestCase(unittest.TestCase): ) assert resp.status_code == 200 + def test_get_status_sig(self): + resp = self.app.get('/status_sig') + assert resp.status_code == 201 + + def test_get_status_response(self): + resp = self.app.get('/status_response') + assert resp.status_code == 201 + def test_serversideerror(self): r = self.app.get('/divide_by_zero') assert r.status_code == 500 diff --git a/wsmeext/flask.py b/wsmeext/flask.py index 33ca52b..f840242 100644 --- a/wsmeext/flask.py +++ b/wsmeext/flask.py @@ -60,14 +60,11 @@ def signature(*args, **kw): try: result = f(*args, **kwargs) - # NOTE: Support setting of status_code with default 200 - status_code = 200 + # NOTE: Support setting of status_code with default 20 + status_code = funcdef.status_code if isinstance(result, wsme.api.Response): - result = result.obj status_code = result.status_code - - if funcdef.status_code: - status_code = funcdef.status_code + result = result.obj res = flask.make_response( dataformat.encode_result(