1 files changed, 17 insertions, 5 deletions

diff --git a/src/browser.rs b/src/browser.rs
index 306bfe1..902f7f7 100644
--- a/src/browser.rs
+++ b/src/browser.rs
@@ -71,6 +71,7 @@ impl LocalBrowser {
         options: FirefoxOptions,
         marionette_port: u16,
         jsdebugger: bool,
+        profile_root: Option<&Path>,
     ) -> WebDriverResult<LocalBrowser> {
         let binary = options.binary.ok_or_else(|| {
             WebDriverError::new(
@@ -87,7 +88,7 @@ impl LocalBrowser {
         let mut profile = match options.profile {
             ProfileType::Named => None,
             ProfileType::Path(x) => Some(x),
-            ProfileType::Temporary => Some(Profile::new()?),
+            ProfileType::Temporary => Some(Profile::new(profile_root)?),
         };
 
         let (profile_path, prefs_backup) = if let Some(ref mut profile) = profile {
@@ -159,10 +160,12 @@ impl LocalBrowser {
             }
         }
         self.process.kill()?;
+
         // Restoring the prefs if the browser fails to stop perhaps doesn't work anyway
         if let Some(prefs_backup) = self.prefs_backup {
             prefs_backup.restore();
         };
+
         Ok(())
     }
 
@@ -227,6 +230,7 @@ fn read_marionette_port(profile_path: &Path) -> Option<u16> {
 pub(crate) struct RemoteBrowser {
     handler: AndroidHandler,
     marionette_port: u16,
+    prefs_backup: Option<PrefsBackup>,
 }
 
 impl RemoteBrowser {
@@ -234,6 +238,7 @@ impl RemoteBrowser {
         options: FirefoxOptions,
         marionette_port: u16,
         websocket_port: Option<u16>,
+        profile_root: Option<&Path>,
     ) -> WebDriverResult<RemoteBrowser> {
         let android_options = options.android.unwrap();
 
@@ -248,10 +253,10 @@ impl RemoteBrowser {
                 ));
             }
             ProfileType::Path(x) => (x, true),
-            ProfileType::Temporary => (Profile::new()?, false),
+            ProfileType::Temporary => (Profile::new(profile_root)?, false),
         };
 
-        set_prefs(
+        let prefs_backup = set_prefs(
             handler.marionette_target_port,
             &mut profile,
             is_custom_profile,
@@ -272,11 +277,18 @@ impl RemoteBrowser {
         Ok(RemoteBrowser {
             handler,
             marionette_port,
+            prefs_backup,
         })
     }
 
     fn close(self) -> WebDriverResult<()> {
         self.handler.force_stop()?;
+
+        // Restoring the prefs if the browser fails to stop perhaps doesn't work anyway
+        if let Some(prefs_backup) = self.prefs_backup {
+            prefs_backup.restore();
+        };
+
         Ok(())
     }
 
@@ -398,7 +410,7 @@ mod tests {
     // several regressions related to remote.log.level.
     #[test]
     fn test_remote_log_level() {
-        let mut profile = Profile::new().unwrap();
+        let mut profile = Profile::new(None).unwrap();
         set_prefs(2828, &mut profile, false, vec![], false).ok();
         let user_prefs = profile.user_prefs().unwrap();
 
@@ -460,7 +472,7 @@ mod tests {
 
     #[test]
     fn test_pref_backup() {
-        let mut profile = Profile::new().unwrap();
+        let mut profile = Profile::new(None).unwrap();
 
         // Create some prefs in the profile
         let initial_prefs = profile.user_prefs().unwrap();