Skip to content

Commit

Permalink
Fix bug that prevents returning Flask.Response's
Browse files Browse the repository at this point in the history
which
https://flask-rebar.readthedocs.io/en/latest/quickstart/basics.html#marshaling
promises should be supported:

  > if response_body_schema is None, the return value must be a return
  > value that Flask supports, e.g. a string or a Flask.Response object
  • Loading branch information
twosigmajab committed Dec 3, 2020
1 parent 9b9a016 commit cf2c567
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 4 deletions.
14 changes: 11 additions & 3 deletions flask_rebar/rebar.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ def _unpack_view_func_return_value(rv):
This imitates Flask's own `Flask.make_response` method.
:param Optional[tuple] rv: (body, status, headers), (body, status), or (body, headers)
:param rv: (body, status, headers), (body, status), (body, headers), or body
:return: (body, status, headers)
:rtype: tuple
"""
Expand Down Expand Up @@ -146,11 +146,19 @@ def wrapped(*args, **kwargs):
if not response_body_schema:
return rv

data, status_code, headers = _unpack_view_func_return_value(rv)
if isinstance(rv, current_app.response_class):
schema = response_body_schema[rv.status_code]
# The schema may be set to None to bypass marshaling (e.g. for 204 responses).
if schema is None:
return rv
# Otherwise, ensure the response body conforms to the promised schema.
schema.loads(rv.data) # May raise ValidationError.
return rv

data, status_code, headers = _unpack_view_func_return_value(rv)
schema = response_body_schema[status_code] # May raise KeyError.

# The schema may be declared as None to bypass marshaling (e.g. for 204 responses).
# The schema may be set to None to bypass marshaling (e.g. for 204 responses).
if schema is None:
return response(
data=data, status_code=status_code, headers=headers, mimetype=mimetype
Expand Down
42 changes: 41 additions & 1 deletion tests/test_rebar.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
import unittest

import marshmallow as m
from flask import Flask
from flask import Flask, make_response

from flask_rebar import messages
from flask_rebar import HeaderApiKeyAuthenticator, SwaggerV3Generator
Expand Down Expand Up @@ -248,6 +248,46 @@ def update_foo(foo_uid):
)
self.assertEqual(resp.status_code, 400)

def test_flask_response_instance_interop_body_matches_schema(self):
rebar = Rebar()
registry = rebar.create_handler_registry()
schema = FooUpdateSchema()

@registry.handles(rule="/foo", method="PUT", response_body_schema=schema)
def foo():
return make_response(({"name": "foo"}, {"foo": "bar"}))

app = create_rebar_app(rebar)
resp = app.test_client().put(path="/foo")
self.assertEqual(resp.status_code, 200)
self.assertEqual(resp.headers["foo"], "bar")

def test_flask_response_instance_interop_body_does_not_match_schema(self):
rebar = Rebar()
registry = rebar.create_handler_registry()
schema = FooUpdateSchema()

@registry.handles(rule="/foo", method="PUT", response_body_schema=schema)
def foo():
return make_response({"does not match": "foo schema"})

app = create_rebar_app(rebar)
resp = app.test_client().put(path="/foo")
self.assertEqual(resp.status_code, 500)

def test_flask_response_instance_interop_no_schema(self):
rebar = Rebar()
registry = rebar.create_handler_registry()

@registry.handles(rule="/foo", response_body_schema={302: None})
def foo():
return make_response(("Redirecting", 302, {"Location": "http://foo.com"}))

app = create_rebar_app(rebar)
resp = app.test_client().get(path="/foo")
self.assertEqual(resp.status_code, 302)
self.assertEqual(resp.headers["Location"], "http://foo.com")

def test_validate_query_parameters(self):
rebar = Rebar()
registry = rebar.create_handler_registry()
Expand Down

0 comments on commit cf2c567

Please sign in to comment.