Python custom decorator in unittest not tearing down - python

In my test case for a web app, I have the need to have different setup/teardowns for my functions, like logging in and auto logging out after logging in, or auto deleting a post after a user posts something. So I build my own decorators like this:
def log_in_and_out(user):
def decorator(f):
#wraps(f)
def decorated_function(self, *args, **kwargs):
try:
userInstance = getattr(self,user)
except AttributeError as e:
e.args = (f"User {user} is not yet defined!",)
raise
self.client.post('/auth/login', data = {
"email": userInstance.email,
"password": "password"}, follow_redirects = True)
f(self)
self.client.get('/auth/logout', follow_redirects = True)
# Fail-safe to ensure user is actually logged out.
# Accesssing "/admin" should redirect to login page.
resp2 = self.client.get('/admin', follow_redirects = True)
self.assertTrue("/auth/login" in getattr(resp2.request,'full_path'))
return decorated_function
return decorator
def post_something(user, post_body = "test"):
def decorator(f):
#wraps(f)
def decorated_function(self, *args, **kwargs):
try:
userInstance = getattr(self,user)
except AttributeError as e:
e.args = (f"User {user} is not yet defined!",)
raise
post = Post(body=post_body, author = userInstance)
db.session.add(post)
db.session.commit()
f(self)
db.session.delete(post)
db.session.commit()
return decorated_function
return decorator
And using them in my test case like this for example (I omit any setupclass/teardownclass for brevity):
class MainTestCase(unittest.TestCase):
#post_something("genericUser","MyNewPost")
#log_in_and_out("genericUser")
def test_post_page(self):
resp = self.client.get("/post/1")
self.assertTrue("MyNewPost" in resp.get_data(as_text=True))
self.assertTrue(f"{self.genericUser.username}" in resp.get_data(as_text = True))
It works if the code runs properly, but if the assertTrue fails, the parts of the decorator under f(self) (e.g. the "teardown" parts) are not run. This messes up the following test cases, since the state of the database are then not reset to a clean state. How can I still make them run? I know if I use pytest instead of unittest, using pytest fixtures solves this issue, but I would like to learn if this is possible to do on my own.

Related

Extracting user id from token through a decorator

I'm working on flask RESTful application that uses auth0 for authentication and authorization. Then, I wrote a decorator that validates the token and extract user's id from it. My goal is use that id, extracted from token, to used inside the decorated function, and throw an exception if user's id from token and from URL parameter doesn't match. This is aimed to avoid users to change data of another user, with his own token. I'm not sure this is the best practice for a RESTful app, but seems to be needed in my case.
That said, I trying to figure out the best approach pass the user's id from token, to the decorated funcion:
Something like this:
def authorization():
def inner(func):
def wrapper(*args, **kwargs):
try:
"""
token validation stuff
...
"""
wrapper.user_id = token_payload['user_id']
except Exception:
return {"success": False}, 500
return func(*args, **kwargs)
return wrapper
return inner
#authorization()
#app.route('/test', methods=['GET'])
def some_function():
return jsonify({
'success': True,
'user_id': some_function.user_id
})
As you can see, I setting the user_id field to the wrapper function, which seems to not be the best way, to do it. Is there any different approach to this situation? maybe using Flask resources?
You can skip one level of wrapping in your decorator since you're not giving it any parameters.
Also, I'd just pass the extracted id into the wrapped function directly instead of setting an attribute.
And finally, you should add the auth decorater innermost, since the decorated function is what you want to register with Flask.
Extra: use functools.wraps to update the signature of your wrapped function, to make introspection and debugging easier.
Thus:
from functools import wraps
def authorized(func):
#wraps(func)
def wrapper(*args, **kw):
try:
# token stuff
user_id = token_payload['user_id']
except:
return jsonify({"success": False}), 500
return func(user_id, *args, **kw)
return wrapper
#app.route('/test', methods=['GET'])
#authorized
def some_function(user_id):
return jsonify({
'success': True,
'user_id': user_id
})
Now every function you decorate with #authorized will need to have user_id as their first parameter, and everything should work as you expect.

create custom decorator for views that check if user is blacklisted

I want to create a decorator for my views that checks the DB if the user has over or equal to 5 warnings, and if so then I redirect him to a page telling him he has been flagged etc.
I wrote the code for it but I want to pass the request object to the decorator so I can call request.user for db queries but that doesn't seem like it's possible. How would I go on about this?
Decorator:
def check_blacklist(request):
try:
db = blacklist.objects.get(user=request.user)
warnings = db.warning
if warnings >= 5:
return redirect("security:flagged")
except:
pass
models.py
class blacklist(models.Model):
user = models.ForeignKey(User, default=None, on_delete=models.CASCADE)
warning = models.IntegerField()
views.py
#check_blacklist(request)
def task_info(request):
return render(request, 'dashboard/task.html',)
I would do something like that :
Use #app.before_request decorator to trap everything before it reachs your router
Define your blacklist fonction to be called before every request
Use request object, as it's global
It should give your something like this (not tested):
#app.before_request
def check_blacklist():
if must_check_blacklist(request.endpoint):
try:
db = blacklist.objects.get(user=request.user)
warnings = db.warning
if warnings >= 5:
return redirect("security:flagged")
except:
pass
I would suggest to call the decorator directly in your urls.py file (it is possible, and I'm doing it for something similar to your usecase). This is what I would do for instance:
# decorators.py
def check_blacklist():
def decorator_fun(func):
def wrapper_func(request, *args, **kwargs):
my_user = request.user
called_func = func(request, *args, **kwargs)
# your stuff goes here
return called_func
return wrapper_func
return decorator_fun
# urls.py
from app.decorators import check_blacklist
urlpatterns = [
path('your_url', check_blacklist()(your_view.as_view()), name='your_view'),
]
Hope that helps!

how to properly refactor this function into another file? django

I have a try and except which is used super often, so I am thinking of taking it out and make it into a function in other file but I cannot seem to make it work.
Can someone please give me a hand?
this is the code which I use really really often
try:
user = get_object_or_404(User, email=data['email'])
except Exception as e:
print(e)
return JSONresponse(False, e.message)
I tried making it into another file and did it this way
def auth_token(data):
try:
return get_object_or_404(User, email=data['email'])
except Exception as e:
return JSONresponse({'status': False})
then when I call auth_token I did it this way in another way and of course I did import it
user = auth_token(data)
I can understand this would then return me JSONresponse({'status': False}) and I tried using raise and it would tell me that I can't use JSONresponse
Can someone give me an idea how this can be done?
Best practice is to use as simple parameters as possible (this makes your function easier to test).
What you're looking for in the exception case is something that can be sent from auth_token through your view:
def auth_token(email): # this returns a User not a token..?
try:
return User.objects.get(email=email)
except User.DoesNotExist:
raise ImmediateHttpResponse(JSONresponse({'status': False}))
Usage would then be:
def myview(request, ...):
token = auth_token(data['email'])
...
Any ImmediateHttpResponse exception will escape the view and must be picked up by a middleware class.
Unfortunately, Django doesn't come with an ImmediateHttpResponse class, but TastyPie has one that we can steal. First the exception class:
class ImmediateHttpResponse(Exception):
"""This exception is used to interrupt the flow of processing to immediately
return a custom HttpResponse.
Common uses include::
* for authentication (like digest/OAuth)
* for throttling
(from TastyPie).
"""
_response = http.HttpResponse("Nothing provided.")
def __init__(self, response):
self._response = response
#property
def response(self):
return self._response
then the middleware:
from myexceptions import ImmediateHttpResponse
class ImmediateHttpResponseMiddleware(object):
"""Middleware that handles ImmediateHttpResponse exceptions, allowing
us to return a response from deep in form handling code.
"""
def process_exception(self, request, exception):
if isinstance(exception, ImmediateHttpResponse):
return exception.response
return None
This class must be added to your MIDDLEWARE_CLASSES in your settings.py file as 'mymiddleware.ImmediateHttpResponseMiddleware'.

Flask-restful - Custom error handling

I want to define custom error handling for a Flask-restful API.
The suggested approach in the documentation here is to do the following:
errors = {
'UserAlreadyExistsError': {
'message': "A user with that username already exists.",
'status': 409,
},
'ResourceDoesNotExist': {
'message': "A resource with that ID no longer exists.",
'status': 410,
'extra': "Any extra information you want.",
},
}
app = Flask(__name__)
api = flask_restful.Api(app, errors=errors)
Now I find this format pretty attractive but I need to specify more parameters when some exception happens. For example, when encountering ResourceDoesNotExist, I want to specify what id does not exist.
Currently, I'm doing the following:
app = Flask(__name__)
api = flask_restful.Api(app)
class APIException(Exception):
def __init__(self, code, message):
self._code = code
self._message = message
#property
def code(self):
return self._code
#property
def message(self):
return self._message
def __str__(self):
return self.__class__.__name__ + ': ' + self.message
class ResourceDoesNotExist(APIException):
"""Custom exception when resource is not found."""
def __init__(self, model_name, id):
message = 'Resource {} {} not found'.format(model_name.title(), id)
super(ResourceNotFound, self).__init__(404, message)
class MyResource(Resource):
def get(self, id):
try:
model = MyModel.get(id)
if not model:
raise ResourceNotFound(MyModel.__name__, id)
except APIException as e:
abort(e.code, str(e))
When called with an id that doesn't exist MyResource will return the following JSON:
{'message': 'ResourceDoesNotExist: Resource MyModel 5 not found'}
This works fine but I'd like to use to Flask-restful error handling instead.
According to the docs
Flask-RESTful will call the handle_error() function on any 400 or 500 error that happens on a Flask-RESTful route, and leave other routes alone.
You can leverage this to implement the required functionality. The only downside is having to create a custom Api.
class CustomApi(flask_restful.Api):
def handle_error(self, e):
flask_restful.abort(e.code, str(e))
If you keep your defined exceptions, when an exception occurs, you'll get the same behaviour as
class MyResource(Resource):
def get(self, id):
try:
model = MyModel.get(id)
if not model:
raise ResourceNotFound(MyModel.__name__, id)
except APIException as e:
abort(e.code, str(e))
Instead of attaching errors dict to Api, I am overriding handle_error method of Api class to handle exceptions of my application.
# File: app.py
# ------------
from flask import Blueprint, jsonify
from flask_restful import Api
from werkzeug.http import HTTP_STATUS_CODES
from werkzeug.exceptions import HTTPException
from view import SomeView
class ExtendedAPI(Api):
"""This class overrides 'handle_error' method of 'Api' class ,
to extend global exception handing functionality of 'flask-restful'.
"""
def handle_error(self, err):
"""It helps preventing writing unnecessary
try/except block though out the application
"""
print(err) # log every exception raised in the application
# Handle HTTPExceptions
if isinstance(err, HTTPException):
return jsonify({
'message': getattr(
err, 'description', HTTP_STATUS_CODES.get(err.code, '')
)
}), err.code
# If msg attribute is not set,
# consider it as Python core exception and
# hide sensitive error info from end user
if not getattr(err, 'message', None):
return jsonify({
'message': 'Server has encountered some error'
}), 500
# Handle application specific custom exceptions
return jsonify(**err.kwargs), err.http_status_code
api_bp = Blueprint('api', __name__)
api = ExtendedAPI(api_bp)
# Routes
api.add_resource(SomeView, '/some_list')
Custom exceptions can be kept in separate file, like:
# File: errors.py
# ---------------
class Error(Exception):
"""Base class for other exceptions"""
def __init__(self, http_status_code:int, *args, **kwargs):
# If the key `msg` is provided, provide the msg string
# to Exception class in order to display
# the msg while raising the exception
self.http_status_code = http_status_code
self.kwargs = kwargs
msg = kwargs.get('msg', kwargs.get('message'))
if msg:
args = (msg,)
super().__init__(args)
self.args = list(args)
for key in kwargs.keys():
setattr(self, key, kwargs[key])
class ValidationError(Error):
"""Should be raised in case of custom validations"""
And in the views exceptions can be raised like:
# File: view.py
# -------------
from flask_restful import Resource
from errors import ValidationError as VE
class SomeView(Resource):
def get(self):
raise VE(
400, # Http Status code
msg='some error message', code=SomeCode
)
Like in view, exceptions can actually be raised from any file in the app which will be handled by the ExtendedAPI handle_error method.
I've used the Blueprint to work with the flask-restful, and I've found that the solution #billmccord and #cedmt provided on the issue not worked for this case, because the Blueprint don't have the handle_exception and handle_user_exception functions.
My workaround is that enhance the function handle_error of the Api, if the "error handler" of the "Exception" have been registered, just raise it, the "error handler" registered on app will deal with that Exception, or the Exception will be passed to the "flask-restful" controlled "custom error handler".
class ImprovedApi(Api):
def handle_error(self, e):
for val in current_app.error_handler_spec.values():
for handler in val.values():
registered_error_handlers = list(filter(lambda x: isinstance(e, x), handler.keys()))
if len(registered_error_handlers) > 0:
raise e
return super().handle_error(e)
api_entry = ImprovedApi(api_entry_bp)
BTW, seems the flask-restful had been deprecated...
I faced the same issue too and after extending flask-restful.Api I realized that
you really don't need to extend the flask-restful.Api
you can easily do this by inheriting from werkzeug.exceptions.HTTPException and solve this issue
app = Flask(__name__)
api = flask_restful.Api(app)
from werkzeug.exceptions import HTTPException
class APIException(HTTPException):
def __init__(self, code, message):
super().__init__()
self.code = code
self.description = message
class ResourceDoesNotExist(APIException):
"""Custom exception when resource is not found."""
def __init__(self, model_name, id):
message = 'Resource {} {} not found'.format(model_name.title(), id)
super().__init__(404, message)
class MyResource(Resource):
def get(self, id):
try:
model = MyModel.get(id)
if not model:
raise ResourceNotFound(MyModel.__name__, id)
except APIException as e:
abort(e.code, str(e))

Developing an error handler as a decorator in flask-restful

I'm trying to develop a rest api by flask-restful. The following decorator is implemented:
def errorhandler(f):
#wraps(f)
def wrapper(*args, **kwargs):
try:
return f(*args, **kwargs)
except errors.DocNotFound:
abort(404, message="Wrong document requested", status=404)
return wrapper
And, Following https://docs.python.org/2/tutorial/errors.html#user-defined-exceptions , in another file named error.py (which is imported here), I have these classes:
class DocError(Exception):
pass
class DocNotFound(DocError):
pass
Now my problem is the implementation of these 2 classes in a way that return my optional error msg. But I don't know how to do that. Would you please give me a hint?
P.S. Here is the way I wanna use the decorator in my Resources:
my_decorators = [
errorhandler,
# Other decorators
]
class Docs(Resource):
method_decorators = my_decorators
def get():
from .errors import DocNotFound
if (<something>):
raise DocNotFound('No access for you!')
return marshal(......)
def delete():
pass
def put():
pass
def post():
pass
Thanks
You can raise your custom exceptions with an argument:
raise DocNotFound('The document "foo" is on the verboten list, no access for you!')
then access that message with:
except errors.DocNotFound as err:
message = err.message or "Wrong document requested"
abort(404, description=message)
The abort(404) call maps to a werkzeug.exceptions.NotFound exception; the description argument lets you override the default description.

Categories