1# coding=UTF-8
2
3# This Source Code Form is subject to the terms of the Mozilla Public
4# License, v. 2.0. If a copy of the MPL was not distributed with this
5# file, You can obtain one at http://mozilla.org/MPL/2.0/.
6
7from __future__ import absolute_import
8
9import os
10import shutil
11import tempfile
12
13import mozprofile
14
15from marionette_driver import errors
16from marionette_harness import MarionetteTestCase
17
18
19class BaseProfileManagement(MarionetteTestCase):
20
21    def setUp(self):
22        super(BaseProfileManagement, self).setUp()
23
24        self.orig_profile_path = self.profile_path
25
26    def tearDown(self):
27        shutil.rmtree(self.orig_profile_path, ignore_errors=True)
28
29        self.marionette.profile = None
30
31        super(BaseProfileManagement, self).tearDown()
32
33    @property
34    def profile(self):
35        return self.marionette.instance.profile
36
37    @property
38    def profile_path(self):
39        return self.marionette.instance.profile.profile
40
41
42class WorkspaceProfileManagement(BaseProfileManagement):
43
44    def setUp(self):
45        super(WorkspaceProfileManagement, self).setUp()
46
47        # Set a new workspace for the instance, which will be used
48        # the next time a new profile is requested by a test.
49        self.workspace = tempfile.mkdtemp()
50        self.marionette.instance.workspace = self.workspace
51
52    def tearDown(self):
53        self.marionette.instance.workspace = None
54
55        shutil.rmtree(self.workspace, ignore_errors=True)
56
57        super(WorkspaceProfileManagement, self).tearDown()
58
59
60class ExternalProfileMixin(object):
61
62    def setUp(self):
63        super(ExternalProfileMixin, self).setUp()
64
65        # Create external profile
66        tmp_dir = tempfile.mkdtemp(suffix="external")
67        shutil.rmtree(tmp_dir, ignore_errors=True)
68
69        self.external_profile = mozprofile.Profile(profile=tmp_dir)
70        # Prevent profile from being removed during cleanup
71        self.external_profile.create_new = False
72
73    def tearDown(self):
74        shutil.rmtree(self.external_profile.profile, ignore_errors=True)
75
76        super(ExternalProfileMixin, self).tearDown()
77
78
79class TestQuitRestartWithoutWorkspace(BaseProfileManagement):
80
81    def test_quit_keeps_same_profile(self):
82        self.marionette.quit()
83        self.marionette.start_session()
84
85        self.assertEqual(self.profile_path, self.orig_profile_path)
86        self.assertTrue(os.path.exists(self.orig_profile_path))
87
88    def test_quit_clean_creates_new_profile(self):
89        self.marionette.quit(clean=True)
90        self.marionette.start_session()
91
92        self.assertNotEqual(self.profile_path, self.orig_profile_path)
93        self.assertFalse(os.path.exists(self.orig_profile_path))
94
95    def test_restart_keeps_same_profile(self):
96        self.marionette.restart()
97
98        self.assertEqual(self.profile_path, self.orig_profile_path)
99        self.assertTrue(os.path.exists(self.orig_profile_path))
100
101    def test_restart_clean_creates_new_profile(self):
102        self.marionette.restart(clean=True)
103
104        self.assertNotEqual(self.profile_path, self.orig_profile_path)
105        self.assertFalse(os.path.exists(self.orig_profile_path))
106
107
108class TestQuitRestartWithWorkspace(WorkspaceProfileManagement):
109
110    def test_quit_keeps_same_profile(self):
111        self.marionette.quit()
112        self.marionette.start_session()
113
114        self.assertEqual(self.profile_path, self.orig_profile_path)
115        self.assertNotIn(self.workspace, self.profile_path)
116        self.assertTrue(os.path.exists(self.orig_profile_path))
117
118    def test_quit_clean_creates_new_profile(self):
119        self.marionette.quit(clean=True)
120        self.marionette.start_session()
121
122        self.assertNotEqual(self.profile_path, self.orig_profile_path)
123        self.assertIn(self.workspace, self.profile_path)
124        self.assertFalse(os.path.exists(self.orig_profile_path))
125
126    def test_restart_keeps_same_profile(self):
127        self.marionette.restart()
128
129        self.assertEqual(self.profile_path, self.orig_profile_path)
130        self.assertNotIn(self.workspace, self.profile_path)
131        self.assertTrue(os.path.exists(self.orig_profile_path))
132
133    def test_restart_clean_creates_new_profile(self):
134        self.marionette.restart(clean=True)
135
136        self.assertNotEqual(self.profile_path, self.orig_profile_path)
137        self.assertIn(self.workspace, self.profile_path)
138        self.assertFalse(os.path.exists(self.orig_profile_path))
139
140
141class TestSwitchProfileFailures(BaseProfileManagement):
142
143    def test_raise_for_switching_profile_while_instance_is_running(self):
144        with self.assertRaisesRegexp(errors.MarionetteException, "instance is not running"):
145            self.marionette.instance.switch_profile()
146
147
148class TestSwitchProfileWithoutWorkspace(ExternalProfileMixin, BaseProfileManagement):
149
150    def setUp(self):
151        super(TestSwitchProfileWithoutWorkspace, self).setUp()
152
153        self.marionette.quit()
154
155    def test_do_not_call_cleanup_of_profile_for_path_only(self):
156        # If a path to a profile has been given (eg. via the --profile command
157        # line argument) and the profile hasn't been created yet, switching the
158        # profile should not try to call `cleanup()` on a string.
159        self.marionette.instance._profile = self.external_profile.profile
160        self.marionette.instance.switch_profile()
161
162    def test_new_random_profile_name(self):
163        self.marionette.instance.switch_profile()
164        self.marionette.start_session()
165
166        self.assertNotEqual(self.profile_path, self.orig_profile_path)
167        self.assertFalse(os.path.exists(self.orig_profile_path))
168
169    def test_new_named_profile(self):
170        self.marionette.instance.switch_profile("foobar")
171        self.marionette.start_session()
172
173        self.assertNotEqual(self.profile_path, self.orig_profile_path)
174        self.assertIn("foobar", self.profile_path)
175        self.assertFalse(os.path.exists(self.orig_profile_path))
176
177    def test_new_named_profile_unicode(self):
178        """Test using unicode string with 1-4 bytes encoding works."""
179        self.marionette.instance.switch_profile(u"$¢€��")
180        self.marionette.start_session()
181
182        self.assertNotEqual(self.profile_path, self.orig_profile_path)
183        self.assertIn(u"$¢€��", self.profile_path)
184        self.assertFalse(os.path.exists(self.orig_profile_path))
185
186    def test_new_named_profile_unicode_escape_characters(self):
187        """Test using escaped unicode string with 1-4 bytes encoding works."""
188        self.marionette.instance.switch_profile(u"\u0024\u00A2\u20AC\u1F36A")
189        self.marionette.start_session()
190
191        self.assertNotEqual(self.profile_path, self.orig_profile_path)
192        self.assertIn(u"\u0024\u00A2\u20AC\u1F36A", self.profile_path)
193        self.assertFalse(os.path.exists(self.orig_profile_path))
194
195    def test_clone_existing_profile(self):
196        self.marionette.instance.switch_profile(clone_from=self.external_profile)
197        self.marionette.start_session()
198
199        self.assertIn(os.path.basename(self.external_profile.profile), self.profile_path)
200        self.assertTrue(os.path.exists(self.external_profile.profile))
201
202    def test_replace_with_current_profile(self):
203        self.marionette.instance.profile = self.profile
204        self.marionette.start_session()
205
206        self.assertEqual(self.profile_path, self.orig_profile_path)
207        self.assertTrue(os.path.exists(self.orig_profile_path))
208
209    def test_replace_with_external_profile(self):
210        self.marionette.instance.profile = self.external_profile
211        self.marionette.start_session()
212
213        self.assertEqual(self.profile_path, self.external_profile.profile)
214        self.assertFalse(os.path.exists(self.orig_profile_path))
215
216        # Set a new profile and ensure the external profile has not been deleted
217        self.marionette.quit()
218        self.marionette.instance.profile = None
219
220        self.assertNotEqual(self.profile_path, self.external_profile.profile)
221        self.assertTrue(os.path.exists(self.external_profile.profile))
222
223
224class TestSwitchProfileWithWorkspace(ExternalProfileMixin, WorkspaceProfileManagement):
225
226    def setUp(self):
227        super(TestSwitchProfileWithWorkspace, self).setUp()
228
229        self.marionette.quit()
230
231    def test_new_random_profile_name(self):
232        self.marionette.instance.switch_profile()
233        self.marionette.start_session()
234
235        self.assertNotEqual(self.profile_path, self.orig_profile_path)
236        self.assertIn(self.workspace, self.profile_path)
237        self.assertFalse(os.path.exists(self.orig_profile_path))
238
239    def test_new_named_profile(self):
240        self.marionette.instance.switch_profile("foobar")
241        self.marionette.start_session()
242
243        self.assertNotEqual(self.profile_path, self.orig_profile_path)
244        self.assertIn("foobar", self.profile_path)
245        self.assertIn(self.workspace, self.profile_path)
246        self.assertFalse(os.path.exists(self.orig_profile_path))
247
248    def test_clone_existing_profile(self):
249        self.marionette.instance.switch_profile(clone_from=self.external_profile)
250        self.marionette.start_session()
251
252        self.assertNotEqual(self.profile_path, self.orig_profile_path)
253        self.assertIn(self.workspace, self.profile_path)
254        self.assertIn(os.path.basename(self.external_profile.profile), self.profile_path)
255        self.assertTrue(os.path.exists(self.external_profile.profile))
256