Bug 1030442 - Be able to set and clear prefs at gecko startup, r=jgriffin

testing/marionette/client/marionette/geckoinstance.py

@@ -2,6 +2,7 @@
 # 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/
 
+from copy import deepcopy
 import os
 import time
 
@@ -21,18 +22,21 @@ class GeckoInstance(object):
                       "browser.warnOnQuit": False}
 
     def __init__(self, host, port, bin, profile, app_args=None, symbols_path=None,
-                  gecko_log=None):
+                  gecko_log=None, prefs=None):
         self.marionette_host = host
         self.marionette_port = port
         self.bin = bin
         self.profile_path = profile
+        self.prefs = prefs
         self.app_args = app_args or []
         self.runner = None
         self.symbols_path = symbols_path
         self.gecko_log = gecko_log
 
     def start(self):
-        profile_args = {"preferences": self.required_prefs}
+        profile_args = {"preferences": deepcopy(self.required_prefs)}
+        if self.prefs:
+            profile_args["preferences"].update(self.prefs)
         if not self.profile_path:
             profile_args["restore"] = False
             profile = Profile(**profile_args)
@@ -75,6 +79,13 @@ class GeckoInstance(object):
             self.runner.stop()
             self.runner.cleanup()
 
+    def restart(self, prefs=None):
+        self.close()
+        if prefs:
+            self.prefs = prefs
+        else:
+            self.prefs = None
+        self.start()
 
 class B2GDesktopInstance(GeckoInstance):
     required_prefs = {"focusmanager.testmode": True}

testing/marionette/client/marionette/marionette.py

@@ -4,6 +4,7 @@
 
 import ConfigParser
 import datetime
+import json
 import os
 import socket
 import StringIO
@@ -688,6 +689,14 @@ class Marionette(object):
                 raise errors.MarionetteException(message=message, status=status, stacktrace=stacktrace)
         raise errors.MarionetteException(message=response, status=500)
 
+    def _reset_timeouts(self):
+        if self.timeout is not None:
+            self.timeouts(self.TIMEOUT_SEARCH, self.timeout)
+            self.timeouts(self.TIMEOUT_SCRIPT, self.timeout)
+            self.timeouts(self.TIMEOUT_PAGE, self.timeout)
+        else:
+            self.timeouts(self.TIMEOUT_PAGE, 30000)
+
     def check_for_crash(self):
         returncode = None
         name = None
@@ -705,6 +714,63 @@ class Marionette(object):
                 (name, returncode))
         return crashed
 
+    def enforce_gecko_prefs(self, prefs):
+        """
+        Checks if the running instance has the given prefs. If not, it will kill the
+        currently running instance, and spawn a new 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 = '%s';
+            let value = '%s';
+            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;
+            }
+            """ % (pref, value))
+            if not pref_exists:
+                break
+        self.set_context(self.CONTEXT_CONTENT)
+        if not pref_exists:
+            self.delete_session()
+            self.instance.restart(prefs)
+            assert(self.wait_for_port()), "Timed out waiting for port!"
+            self.start_session()
+            self._reset_timeouts()
+
+    def restart_with_clean_profile(self):
+        """
+        This will terminate the currently running instance, and spawn a new instance
+        with a clean profile.
+
+        : 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")
+        self.delete_session()
+        self.instance.restart()
+        assert(self.wait_for_port()), "Timed out waiting for port!"
+        self.start_session()
+        self._reset_timeouts()
+
     def absolute_url(self, relative_url):
         '''
         Returns an absolute url for files served from Marionette's www directory.

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

@@ -0,0 +1,32 @@
+# 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 os
+from errors import JavascriptException
+from marionette_test import MarionetteTestCase
+
+class TestLog(MarionetteTestCase):
+    def setUp(self):
+        MarionetteTestCase.setUp(self)
+        self.marionette.enforce_gecko_prefs({"marionette.test.bool": True, "marionette.test.string": "testing", "marionette.test.int": 3})
+
+    def test_preferences_are_set(self):
+        bool_value = self.marionette.execute_script("return SpecialPowers.getBoolPref('marionette.test.bool');")
+        string_value = self.marionette.execute_script("return SpecialPowers.getCharPref('marionette.test.string');")
+        int_value = self.marionette.execute_script("return SpecialPowers.getIntPref('marionette.test.int');")
+        self.assertTrue(bool_value)
+        self.assertEqual(string_value, "testing")
+        self.assertEqual(int_value, 3)
+
+    def test_change_preset(self):
+        bool_value = self.marionette.execute_script("return SpecialPowers.getBoolPref('marionette.test.bool');")
+        self.assertTrue(bool_value)
+        self.marionette.enforce_gecko_prefs({"marionette.test.bool": False})
+        bool_value = self.marionette.execute_script("return SpecialPowers.getBoolPref('marionette.test.bool');")
+        self.assertFalse(bool_value)
+
+    def test_clean_profile(self):
+        self.marionette.restart_with_clean_profile()
+        with self.assertRaisesRegexp(JavascriptException, "Error getting pref"):
+            bool_value = self.marionette.execute_script("return SpecialPowers.getBoolPref('marionette.test.bool');")

testing/marionette/client/marionette/tests/unit/unit-tests.ini

@@ -115,5 +115,7 @@ browser = false
 
 [test_execute_isolate.py]
 [test_click_scrolling.py]
+[test_profile_management.py]
+b2g = false
 
 [include:oop/manifest.ini]