--- a/testing/marionette/client/marionette_driver/marionette.py
+++ b/testing/marionette/client/marionette_driver/marionette.py
@@ -1043,47 +1043,51 @@ class Marionette(object):
         instance with the requested preferences.
 
         : param prefs: A dictionary whose keys are preference names.
         """
         if not self.instance:
             raise errors.MarionetteException("enforce_gecko_prefs() can only be called "
                                              "on Gecko instances launched by Marionette")
         pref_exists = True
-        self.set_context(self.CONTEXT_CHROME)
-        for pref, value in prefs.iteritems():
-            if type(value) is not str:
-                value = json.dumps(value)
-            pref_exists = self.execute_script("""
-            let prefInterface = Components.classes["@mozilla.org/preferences-service;1"]
-                                          .getService(Components.interfaces.nsIPrefBranch);
-            let pref = '{0}';
-            let value = '{1}';
-            let type = prefInterface.getPrefType(pref);
-            switch(type) {{
-                case prefInterface.PREF_STRING:
-                    return value == prefInterface.getCharPref(pref).toString();
-                case prefInterface.PREF_BOOL:
-                    return value == prefInterface.getBoolPref(pref).toString();
-                case prefInterface.PREF_INT:
-                    return value == prefInterface.getIntPref(pref).toString();
-                case prefInterface.PREF_INVALID:
-                    return false;
-            }}
-            """.format(pref, value))
-            if not pref_exists:
-                break
-        self.set_context(self.CONTEXT_CONTENT)
+        with self.using_context(self.CONTEXT_CHROME):
+            for pref, value in prefs.iteritems():
+                if type(value) is not str:
+                    value = json.dumps(value)
+                pref_exists = self.execute_script("""
+                let prefInterface = Components.classes["@mozilla.org/preferences-service;1"]
+                                              .getService(Components.interfaces.nsIPrefBranch);
+                let pref = '{0}';
+                let value = '{1}';
+                let type = prefInterface.getPrefType(pref);
+                switch(type) {{
+                    case prefInterface.PREF_STRING:
+                        return value == prefInterface.getCharPref(pref).toString();
+                    case prefInterface.PREF_BOOL:
+                        return value == prefInterface.getBoolPref(pref).toString();
+                    case prefInterface.PREF_INT:
+                        return value == prefInterface.getIntPref(pref).toString();
+                    case prefInterface.PREF_INVALID:
+                        return false;
+                }}
+                """.format(pref, value))
+                if not pref_exists:
+                    break
+
         if not pref_exists:
+            context = self._send_message("getContext", key="value")
             self.delete_session()
             self.instance.restart(prefs)
             self.raise_for_port()
             self.start_session()
             self.reset_timeouts()
 
+            # Restore the context as used before the restart
+            self.set_context(context)
+
     def _request_in_app_shutdown(self, shutdown_flags=None):
         """Terminate the currently running instance from inside the application.
 
         :param shutdown_flags: If specified use additional flags for the shutdown
                                of the application. Possible values here correspond
                                to constants in nsIAppStartup: http://mzl.la/1X0JZsC.
         """
         flags = set([])
@@ -1152,16 +1156,18 @@ class Marionette(object):
                        browser. Otherwise the browser will be restarted immediately
                        by killing the process.
         :param callback: If provided and `in_app` is True, the callback will be
                          used to trigger the restart.
         """
         if not self.instance:
             raise errors.MarionetteException("restart() can only be called "
                                              "on Gecko instances launched by Marionette")
+
+        context = self._send_message("getContext", key="value")
         session_id = self.session_id
 
         if in_app:
             if clean:
                 raise ValueError("An in_app restart cannot be triggered with the clean flag set")
 
             if callable(callback):
                 self._send_message("acceptConnections", {"value": False})
@@ -1181,16 +1187,19 @@ class Marionette(object):
         else:
             self.delete_session()
             self.instance.restart(clean=clean)
             self.raise_for_port()
 
         self.start_session(session_id=session_id)
         self.reset_timeouts()
 
+        # Restore the context as used before the restart
+        self.set_context(context)
+
         if in_app and self.session.get("processId"):
             # In some cases Firefox restarts itself by spawning into a new process group.
             # As long as mozprocess cannot track that behavior (bug 1284864) we assist by
             # informing about the new process id.
             self.instance.runner.process_handler.check_for_detached(self.session["processId"])
 
     def absolute_url(self, relative_url):
         '''

--- a/testing/marionette/harness/marionette/tests/unit/test_profile_management.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_profile_management.py
@@ -20,17 +20,15 @@ class TestProfileManagement(MarionetteTe
         self.assertTrue(self.marionette.get_pref("marionette.test.bool"))
         self.assertEqual(self.marionette.get_pref("marionette.test.string"), "testing")
         self.assertEqual(self.marionette.get_pref("marionette.test.int"), 3)
 
     def test_change_preference(self):
         self.assertTrue(self.marionette.get_pref("marionette.test.bool"))
 
         self.marionette.enforce_gecko_prefs({"marionette.test.bool": False})
-        self.marionette.set_context('chrome')
 
         self.assertFalse(self.marionette.get_pref("marionette.test.bool"))
 
     def test_clean_profile(self):
         self.marionette.restart(clean=True)
-        self.marionette.set_context('chrome')
 
         self.assertEqual(self.marionette.get_pref("marionette.test.bool"), None)

--- a/testing/marionette/harness/marionette/tests/unit/test_quit_restart.py
+++ b/testing/marionette/harness/marionette/tests/unit/test_quit_restart.py
@@ -95,16 +95,76 @@ class TestQuitRestart(MarionetteTestCase
         self.assertEqual(self.marionette.session, None)
         with self.assertRaisesRegexp(MarionetteException, "Please start a session"):
             self.marionette.get_url()
 
         self.marionette.start_session()
         self.assertNotEqual(self.marionette.session_id, self.session_id)
         self.assertNotEqual(self.marionette.get_pref("browser.startup.page"), 3)
 
+    def test_reset_context_after_quit_by_set_context(self):
+        # Check that we are in content context which is used by default in Marionette
+        self.assertNotIn('chrome://', self.marionette.get_url(),
+                         "Context doesn't default to content")
+
+        self.marionette.set_context('chrome')
+        self.marionette.quit(in_app=True)
+        self.assertEqual(self.marionette.session, None)
+        self.marionette.start_session()
+        self.assertNotIn('chrome://', self.marionette.get_url(),
+                         "Not in content context after quit with using_context")
+
+    def test_reset_context_after_quit_by_using_context(self):
+        # Check that we are in content context which is used by default in Marionette
+        self.assertNotIn('chrome://', self.marionette.get_url(),
+                         "Context doesn't default to content")
+
+        with self.marionette.using_context('chrome'):
+            self.marionette.quit(in_app=True)
+            self.assertEqual(self.marionette.session, None)
+            self.marionette.start_session()
+            self.assertNotIn('chrome://', self.marionette.get_url(),
+                             "Not in content context after quit with using_context")
+
+    def test_keep_context_after_restart_by_set_context(self):
+        # Check that we are in content context which is used by default in Marionette
+        self.assertNotIn('chrome://', self.marionette.get_url(),
+                         "Context doesn't default to content")
+
+        # restart while we are in chrome context
+        self.marionette.set_context('chrome')
+        self.marionette.restart(in_app=True)
+
+        # An in-app restart will keep the same process id only on Linux
+        if self.marionette.session_capabilities['platformName'] == 'linux':
+            self.assertEqual(self.marionette.session["processId"], self.pid)
+        else:
+            self.assertNotEqual(self.marionette.session["processId"], self.pid)
+
+        self.assertIn('chrome://', self.marionette.get_url(),
+                      "Not in chrome context after a restart with set_context")
+
+    def test_keep_context_after_restart_by_using_context(self):
+        # Check that we are in content context which is used by default in Marionette
+        self.assertNotIn('chrome://', self.marionette.get_url(),
+                         "Context doesn't default to content")
+
+        # restart while we are in chrome context
+        with self.marionette.using_context('chrome'):
+            self.marionette.restart(in_app=True)
+
+            # An in-app restart will keep the same process id only on Linux
+            if self.marionette.session_capabilities['platformName'] == 'linux':
+                self.assertEqual(self.marionette.session["processId"], self.pid)
+            else:
+                self.assertNotEqual(self.marionette.session["processId"], self.pid)
+
+            self.assertIn('chrome://', self.marionette.get_url(),
+                          "Not in chrome context after a restart with using_context")
+
     def shutdown(self, restart=False):
         self.marionette.set_context("chrome")
         self.marionette.execute_script("""
             Components.utils.import("resource://gre/modules/Services.jsm");
             let flags = Ci.nsIAppStartup.eAttemptQuit
             if(arguments[0]) {
               flags |= Ci.nsIAppStartup.eRestart;
             }

--- a/testing/puppeteer/firefox/firefox_puppeteer/testcases/base.py
+++ b/testing/puppeteer/firefox/firefox_puppeteer/testcases/base.py
@@ -71,19 +71,16 @@ class BaseFirefoxTestCase(unittest.TestC
 
         :param flags: Specific restart flags for Firefox
         """
         if kwargs.get('clean'):
             self.marionette.restart(clean=True)
         else:
             self.marionette.restart(in_app=True)
 
-        # Marionette doesn't keep the former context, so restore to chrome
-        self.marionette.set_context('chrome')
-
         # Ensure that we always have a valid browser instance available
         self.browser = self.windows.switch_to(lambda win: type(win) is BrowserWindow)
 
     def setUp(self, *args, **kwargs):
         super(BaseFirefoxTestCase, self).setUp(*args, **kwargs)
 
         self._start_handle_count = len(self.marionette.window_handles)
         self._init_tab_handles = self.marionette.window_handles