Bug 1237396 - Add safebrowsing test for initial download of files. r=whimboo

More adjustment on the if statement within the base class.

MozReview-Commit-ID: HjifORUSWXo

--HG--
extra : rebase_source : dc0b10274613b8cbae38086571f4ab52750a2266

testing/firefox-ui/tests/functional/security/manifest.ini

@@ -9,6 +9,7 @@ tags = local
 [test_mixed_script_content_blocking.py]
 [test_no_certificate.py]
 tags = local
+[test_safe_browsing_initial_download.py]
 [test_safe_browsing_notification.py]
 [test_safe_browsing_warning_pages.py]
 [test_security_notification.py]

testing/firefox-ui/tests/functional/security/test_safe_browsing_initial_download.py

@@ -0,0 +1,83 @@
+# 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 firefox_ui_harness.testcases import FirefoxTestCase
+from marionette_driver import Wait
+
+
+class TestSafeBrowsingInitialDownload(FirefoxTestCase):
+
+    test_data = [{
+            'platforms': ['linux', 'windows_nt', 'darwin'],
+            'files': [
+                # Phishing
+                "goog-badbinurl-shavar.cache",
+                "goog-badbinurl-shavar.pset",
+                "goog-badbinurl-shavar.sbstore",
+                "goog-malware-shavar.cache",
+                "goog-malware-shavar.pset",
+                "goog-malware-shavar.sbstore",
+                "goog-phish-shavar.cache",
+                "goog-phish-shavar.pset",
+                "goog-phish-shavar.sbstore",
+                "goog-unwanted-shavar.cache",
+                "goog-unwanted-shavar.pset",
+                "goog-unwanted-shavar.sbstore",
+
+                # Tracking Protections
+                "mozstd-track-digest256.cache",
+                "mozstd-track-digest256.pset",
+                "mozstd-track-digest256.sbstore",
+                "mozstd-trackwhite-digest256.cache",
+                "mozstd-trackwhite-digest256.pset",
+                "mozstd-trackwhite-digest256.sbstore"
+                ]
+            },
+        {
+            'platforms': ['windows_nt'],
+            'files': [
+                "goog-downloadwhite-digest256.cache",
+                "goog-downloadwhite-digest256.pset",
+                "goog-downloadwhite-digest256.sbstore"
+            ]
+        }
+    ]
+
+    browser_prefs = {
+        'browser.safebrowsing.downloads.enabled': 'true',
+        'browser.safebrowsing.enabled': 'true',
+        'browser.safebrowsing.malware.enabled': 'true',
+        'browser.safebrowsing.provider.google.nextupdatetime': 1,
+        'browser.safebrowsing.provider.mozilla.nextupdatetime': 1,
+        'privacy.trackingprotection.enabled': 'true',
+        'privacy.trackingprotection.pbmode.enabled': 'true',
+    }
+
+    def setUp(self):
+        FirefoxTestCase.setUp(self)
+
+        # Set Browser Preferences
+        self.marionette.enforce_gecko_prefs(self.browser_prefs)
+
+        # Get safebrowsing path where downloaded data gets stored
+        self.sb_files_path = os.path.join(self.marionette.instance.profile.profile, 'safebrowsing')
+
+    def tearDown(self):
+        try:
+            self.restart(clean=True)
+        finally:
+            FirefoxTestCase.tearDown(self)
+
+    def test_safe_browsing_initial_download(self):
+        wait = Wait(self.marionette, timeout=self.browser.timeout_page_load)
+
+        for data in self.test_data:
+            if self.platform not in data['platforms']:
+                continue
+            for item in data['files']:
+                wait.until(
+                    lambda _: os.path.exists(os.path.join(self.sb_files_path, item)),
+                    message='Safe Browsing File: {} not found!'.format(item))

testing/puppeteer/firefox/firefox_puppeteer/testcases/base.py

@@ -66,20 +66,25 @@ class BaseFirefoxTestCase(unittest.TestCase, Puppeteer):
             self.browser.tabbar.close_all_tabs([self.browser.tabbar.tabs[0]])
             self.browser.tabbar.tabs[0].switch_to()
 
-    def restart(self, flags=None):
+    def restart(self, **kwargs):
         """Restart Firefox and re-initialize data.
 
         :param flags: Specific restart flags for Firefox
         """
+
         # TODO: Bug 1148220 Marionette's in_app restart has to send 'quit-application-requested'
         # observer notification before an in_app restart
-        self.marionette.execute_script("""
-          Components.utils.import("resource://gre/modules/Services.jsm");
-          let cancelQuit = Components.classes["@mozilla.org/supports-PRBool;1"]
-                                     .createInstance(Components.interfaces.nsISupportsPRBool);
-          Services.obs.notifyObservers(cancelQuit, "quit-application-requested", null);
-        """)
-        self.marionette.restart(in_app=True)
+        with self.marionette.using_context('chrome'):
+            self.marionette.execute_script("""
+                Components.utils.import("resource://gre/modules/Services.jsm");
+                let cancelQuit = Components.classes["@mozilla.org/supports-PRBool;1"]
+                                         .createInstance(Components.interfaces.nsISupportsPRBool);
+                Services.obs.notifyObservers(cancelQuit, "quit-application-requested", null);
+                """)
+        if kwargs.get('clean'):
+            self.marionette.restart(clean=True)
+        else:
+            self.marionette.restart(in_app=True)
 
         # Marionette doesn't keep the former context, so restore to chrome
         self.marionette.set_context('chrome')