diff --git a/src/pypnusershub/tests/utils.py b/src/pypnusershub/tests/utils.py index 9c0bfa9..3f7a72d 100644 --- a/src/pypnusershub/tests/utils.py +++ b/src/pypnusershub/tests/utils.py @@ -1,3 +1,5 @@ +from contextlib import contextmanager + from flask import current_app from werkzeug.http import dump_cookie from werkzeug.datastructures import Headers @@ -30,6 +32,19 @@ def unset_logged_user(client): client.environ_base.pop("HTTP_AUTHORIZATION") +@contextmanager +def logged_user(client, user): + """ + Usage: + + with logged_user(client, user): + response = client.get(url) + """ + set_logged_user(client, user) + yield + unset_logged_user(client) + + # retro compatibility for cookie auth unset_logged_user_cookie = unset_logged_user