Bug 1097676 - Create a convenience decorator for marionette.using_context(), r=ato

testing/marionette/client/marionette/tests/unit/test_with_using_context.py

@@ -5,6 +5,7 @@
 from marionette import MarionetteTestCase
 from marionette_driver.errors import MarionetteException
 
+from marionette_driver.decorators import using_context
 
 class TestSetContext(MarionetteTestCase):
     def setUp(self):
@@ -52,3 +53,13 @@ class TestSetContext(MarionetteTestCase):
             with self.marionette.using_context(self.chrome):
                 raise MarionetteException
         self.assertEquals(self.get_context(), self.content)
+
+    def test_with_using_context_decorator(self):
+        @using_context('content')
+        def inner_content(m):
+            self.assertEquals(self.get_context(), 'content')
+        @using_context('chrome')
+        def inner_chrome(m):
+            self.assertEquals(self.get_context(), 'chrome')
+        inner_content(self.marionette)
+        inner_chrome(self.marionette)

testing/marionette/driver/marionette_driver/decorators.py

@@ -65,3 +65,19 @@ def uses_marionette(func):
 
         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