diff options
Diffstat (limited to 'testing/marionette/client/marionette_driver/decorators.py')
-rw-r--r-- | testing/marionette/client/marionette_driver/decorators.py | 79 |
1 files changed, 79 insertions, 0 deletions
diff --git a/testing/marionette/client/marionette_driver/decorators.py b/testing/marionette/client/marionette_driver/decorators.py new file mode 100644 index 0000000000..95a5c5bbee --- /dev/null +++ b/testing/marionette/client/marionette_driver/decorators.py @@ -0,0 +1,79 @@ +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + +import socket +from functools import wraps + + +def _find_marionette_in_args(*args, **kwargs): + try: + m = [a for a in args + tuple(kwargs.values()) if hasattr(a, "session")][0] + except IndexError: + print("Can only apply decorator to function using a marionette object") + raise + return m + + +def do_process_check(func): + """Decorator which checks the process status after the function has run.""" + + @wraps(func) + def _(*args, **kwargs): + try: + return func(*args, **kwargs) + except (socket.error, socket.timeout): + m = _find_marionette_in_args(*args, **kwargs) + + # In case of socket failures which will also include crashes of the + # application, make sure to handle those correctly. In case of an + # active shutdown just let it bubble up. + if m.is_shutting_down: + raise + + m._handle_socket_failure() + + return _ + + +def uses_marionette(func): + """Decorator which creates a marionette session and deletes it + afterwards if one doesn't already exist. + """ + + @wraps(func) + def _(*args, **kwargs): + m = _find_marionette_in_args(*args, **kwargs) + delete_session = False + if not m.session: + delete_session = True + m.start_session() + + m.set_context(m.CONTEXT_CHROME) + ret = func(*args, **kwargs) + + if delete_session: + m.delete_session() + + return ret + + return _ + + +def using_context(context): + """Decorator which allows a function to execute in certain scope + using marionette.using_context functionality and returns to old + scope once the function exits. + :param context: Either 'chrome' or 'content'. + """ + + def wrap(func): + @wraps(func) + def inner(*args, **kwargs): + m = _find_marionette_in_args(*args, **kwargs) + with m.using_context(context): + return func(*args, **kwargs) + + return inner + + return wrap |