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