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