SeleniumHQ
diff --git a/‎py/selenium/webdriver/firefox/remote_connection.py‎
Lines changed: 1 addition & 0 deletions b/‎py/selenium/webdriver/firefox/remote_connection.py‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎py/selenium/webdriver/firefox/webdriver.py‎
Lines changed: 28 additions & 0 deletions b/‎py/selenium/webdriver/firefox/webdriver.py‎
Lines changed: 28 additions & 0 deletions
diff --git a/‎py/test/selenium/webdriver/marionette/mn_context_tests.py‎
Lines changed: 30 additions & 0 deletions b/‎py/test/selenium/webdriver/marionette/mn_context_tests.py‎
Lines changed: 30 additions & 0 deletions
@@ -22,6 +22,7 @@ class FirefoxRemoteConnection(RemoteConnection):
     def __init__(self, remote_server_addr, keep_alive=True):
         RemoteConnection.__init__(self, remote_server_addr, keep_alive)
 
+        self._commands["GET_CONTEXT"] = ('GET', '/session/$sessionId/moz/context')
         self._commands["SET_CONTEXT"] = ("POST", "/session/$sessionId/moz/context")
         self._commands["ELEMENT_GET_ANONYMOUS_CHILDREN"] = \
             ("POST", "/session/$sessionId/moz/xbl/$id/anonymous_children")
 
@@ -31,6 +31,8 @@
 import types
 
 from .extension_connection import ExtensionConnection
+from contextlib import contextmanager
+
 from .firefox_binary import FirefoxBinary
 from .firefox_profile import FirefoxProfile
 from .options import Options
@@ -130,6 +132,10 @@ def __init__(self, firefox_profile=None, firefox_binary=None,
 
         # W3C remote
         # TODO(ato): Perform conformance negotiation
+
+        self.CONTEXT_CHROME = 'chrome'
+        self.CONTEXT_CONTENT = 'content'
+
         if capabilities.get("marionette"):
             self.service = Service(executable_path, log_path=log_path)
             self.service.start()
@@ -195,3 +201,25 @@ def firefox_profile(self):
 
     def set_context(self, context):
         self.execute("SET_CONTEXT", {"context": context})
+
+    @contextmanager
+    def context(self, context):
+        """Sets the context that Selenium commands are running in using
+        a `with` statement. The state of the context on the server is
+        saved before entering the block, and restored upon exiting it.
+
+        :param context: Context, may be one of the class properties
+            `CONTEXT_CHROME` or `CONTEXT_CONTENT`.
+
+        Usage example::
+
+            with selenium.context(selenium.CONTEXT_CHROME):
+                # chrome scope
+                ... do stuff ...
+        """
+        initial_context = self.execute('GET_CONTEXT').pop('value')
+        self.set_context(context)
+        try:
+            yield
+        finally:
+            self.set_context(initial_context)
@@ -0,0 +1,30 @@
+# Licensed to the Software Freedom Conservancy (SFC) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The SFC licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+class TestUsingContext(object):
+
+    def test_context_sets_correct_context_and_returns(self, driver):
+
+        def get_context():
+            return driver.execute('GET_CONTEXT').pop('value')
+
+        assert get_context() == driver.CONTEXT_CONTENT
+        with driver.context(driver.CONTEXT_CHROME):
+            assert get_context() == driver.CONTEXT_CHROME
+        assert get_context() == driver.CONTEXT_CONTENT
+