1# This Source Code Form is subject to the terms of the Mozilla Public 2# License, v. 2.0. If a copy of the MPL was not distributed with this 3# file, You can obtain one at http://mozilla.org/MPL/2.0/. 4 5from __future__ import absolute_import, print_function 6 7import socket 8 9from functools import wraps 10 11 12def _find_marionette_in_args(*args, **kwargs): 13 try: 14 m = [a for a in args + tuple(kwargs.values()) if hasattr(a, "session")][0] 15 except IndexError: 16 print("Can only apply decorator to function using a marionette object") 17 raise 18 return m 19 20 21def do_process_check(func): 22 """Decorator which checks the process status after the function has run.""" 23 24 @wraps(func) 25 def _(*args, **kwargs): 26 try: 27 return func(*args, **kwargs) 28 except (socket.error, socket.timeout): 29 m = _find_marionette_in_args(*args, **kwargs) 30 31 # In case of socket failures which will also include crashes of the 32 # application, make sure to handle those correctly. In case of an 33 # active shutdown just let it bubble up. 34 if m.is_shutting_down: 35 raise 36 37 m._handle_socket_failure() 38 39 return _ 40 41 42def uses_marionette(func): 43 """Decorator which creates a marionette session and deletes it 44 afterwards if one doesn't already exist. 45 """ 46 47 @wraps(func) 48 def _(*args, **kwargs): 49 m = _find_marionette_in_args(*args, **kwargs) 50 delete_session = False 51 if not m.session: 52 delete_session = True 53 m.start_session() 54 55 m.set_context(m.CONTEXT_CHROME) 56 ret = func(*args, **kwargs) 57 58 if delete_session: 59 m.delete_session() 60 61 return ret 62 63 return _ 64 65 66def using_context(context): 67 """Decorator which allows a function to execute in certain scope 68 using marionette.using_context functionality and returns to old 69 scope once the function exits. 70 :param context: Either 'chrome' or 'content'. 71 """ 72 73 def wrap(func): 74 @wraps(func) 75 def inner(*args, **kwargs): 76 m = _find_marionette_in_args(*args, **kwargs) 77 with m.using_context(context): 78 return func(*args, **kwargs) 79 80 return inner 81 82 return wrap 83