Bug 1521754 - Add marionette test coverage for new private browsing UI with search r=nanj

MozReview-Commit-ID: DbuVFNLew6V

Differential Revision: https://phabricator.services.mozilla.com/D17931

--HG--
extra : moz-landing-system : lando

testing/firefox-ui/tests/functional/private_browsing/test_about_private_browsing.py

@@ -16,16 +16,18 @@ class TestAboutPrivateBrowsing(PuppeteerMixin, MarionetteTestCase):
         # Use a fake local support URL
         support_url = 'about:blank?'
         self.marionette.set_pref('app.support.baseURL', support_url)
-        self.marionette.set_pref('browser.privatebrowsing.searchUI', False)
-
         self.pb_url = support_url + 'private-browsing'
 
+        # Disable the search UI
+        self.marionette.set_pref('browser.privatebrowsing.searchUI', False)
+
     def tearDown(self):
         try:
             self.puppeteer.windows.close_all([self.browser])
             self.browser.switch_to()
 
             self.marionette.clear_pref('app.support.baseURL')
+            self.marionette.clear_pref('browser.privatebrowsing.searchUI')
         finally:
             super(TestAboutPrivateBrowsing, self).tearDown()
 
@@ -62,3 +64,70 @@ class TestAboutPrivateBrowsing(PuppeteerMixin, MarionetteTestCase):
 
         finally:
             pb_window.close()
+
+
+class TestAboutPrivateBrowsingWithSearch(PuppeteerMixin, MarionetteTestCase):
+
+    def setUp(self):
+        super(TestAboutPrivateBrowsingWithSearch, self).setUp()
+
+        # Use a fake local support URL
+        support_url = 'about:blank?'
+        self.marionette.set_pref('app.support.baseURL', support_url)
+        self.pb_url = support_url + 'private-browsing-myths'
+
+        # Enable the search UI
+        self.marionette.set_pref('browser.privatebrowsing.searchUI', True)
+
+    def tearDown(self):
+        try:
+            self.puppeteer.windows.close_all([self.browser])
+            self.browser.switch_to()
+
+            self.marionette.clear_pref('app.support.baseURL')
+            self.marionette.clear_pref('browser.privatebrowsing.searchUI')
+        finally:
+            super(TestAboutPrivateBrowsingWithSearch, self).tearDown()
+
+    def testCheckAboutPrivateBrowsingWithSearch(self):
+        self.assertFalse(self.browser.is_private)
+
+        with self.marionette.using_context('content'):
+            self.marionette.navigate('about:privatebrowsing')
+
+            status_node = self.marionette.find_element(By.CSS_SELECTOR, 'p.showNormal')
+            self.assertEqual(status_node.text,
+                             self.browser.localize_entity('aboutPrivateBrowsing.notPrivate'),
+                             'Status text indicates we are not in private browsing mode')
+
+        def window_opener(win):
+            with win.marionette.using_context('content'):
+                button = self.marionette.find_element(By.ID, 'startPrivateBrowsing')
+                button.click()
+
+        pb_window = self.browser.open_window(callback=window_opener,
+                                             expected_window_class=BrowserWindow)
+
+        try:
+            self.assertTrue(pb_window.is_private)
+
+            # Test the search hand-off
+            with self.marionette.using_context('content'):
+                search = self.marionette.find_element(By.ID, 'search-handoff-button')
+                search.click()
+
+            self.assertTrue(pb_window.navbar.locationbar.focused, 'url bar is focused')
+
+            pb_window.navbar.locationbar.urlbar.send_keys('foo')
+            self.assertEqual(pb_window.navbar.locationbar.value, '@google foo',
+                             'url bar prepends the @search shortcut')
+
+            # Test the private browsing myths link
+            with self.marionette.using_context('content'):
+                link = self.marionette.find_element(By.ID, 'private-browsing-myths')
+                link.click()
+                Wait(self.marionette, timeout=self.marionette.timeout.page_load).until(
+                    lambda _: self.marionette.get_url() == self.pb_url)
+
+        finally:
+            pb_window.close()