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