--- a/testing/firefox-ui/tests/functional/security/test_safe_browsing_initial_download.py
+++ b/testing/firefox-ui/tests/functional/security/test_safe_browsing_initial_download.py
@@ -1,85 +1,80 @@
 # 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
-import re
 
 from firefox_ui_harness.testcases import FirefoxTestCase
 from marionette_driver import Wait
-from marionette_driver.errors import TimeoutException
 
 
 class TestSafeBrowsingInitialDownload(FirefoxTestCase):
 
-    test_data = [{
-        'platforms': ['linux', 'windows_nt', 'darwin'],
-        'files': [
-            # Phishing
-            r'^goog-badbinurl-shavar.pset$',
-            r'^goog-badbinurl-shavar.sbstore$',
-            r'^goog-malware-shavar.pset$',
-            r'^goog-malware-shavar.sbstore$',
-            r'^goog(pub)?-phish-shavar.pset$',
-            r'^goog(pub)?-phish-shavar.sbstore$',
-            r'^goog-unwanted-shavar.pset$',
-            r'^goog-unwanted-shavar.sbstore$',
+    file_extensions = [
+        'pset',
+        'sbstore',
+    ]
 
-            # Tracking Protections
-            r'^base-track-digest256.pset$',
-            r'^base-track-digest256.sbstore$',
-            r'^mozstd-trackwhite-digest256.pset$',
-            r'^mozstd-trackwhite-digest256.sbstore$',
-        ],
-    }, {
-        'platforms': ['windows_nt'],
-        'files': [
-            r'^goog-downloadwhite-digest256.pset$',
-            r'^goog-downloadwhite-digest256.sbstore$',
-        ]
-    },
+    prefs_download_lists = [
+        'urlclassifier.blockedTable',
+        'urlclassifier.downloadAllowTable',
+        'urlclassifier.downloadBlockTable',
+        'urlclassifier.malwareTable',
+        'urlclassifier.phishTable',
+        'urlclassifier.trackingTable',
+        'urlclassifier.trackingWhitelistTable',
     ]
 
-    browser_prefs = {
-        'browser.safebrowsing.downloads.enabled': 'true',
-        'browser.safebrowsing.phishing.enabled': 'true',
-        'browser.safebrowsing.malware.enabled': 'true',
+    prefs_provider_update_time = {
+        # Force an immediate download of the safebrowsing files
         'browser.safebrowsing.provider.google.nextupdatetime': 1,
         'browser.safebrowsing.provider.mozilla.nextupdatetime': 1,
-        'privacy.trackingprotection.enabled': 'true',
-        'privacy.trackingprotection.pbmode.enabled': 'true',
+    }
+
+    prefs_safebrowsing = {
+        'browser.safebrowsing.blockedURIs.enabled': True,
+        'browser.safebrowsing.downloads.enabled': True,
+        'browser.safebrowsing.phishing.enabled': True,
+        'browser.safebrowsing.malware.enabled': True,
+        'privacy.trackingprotection.enabled': True,
+        'privacy.trackingprotection.pbmode.enabled': True,
     }
 
+    def get_safebrowsing_files(self):
+        files = []
+        for pref_name in self.prefs_download_lists:
+            base_names = self.marionette.get_pref(pref_name).split(',')
+            for ext in self.file_extensions:
+                files.extend(['{file}.{ext}'.format(file=f, ext=ext) for f in base_names if f])
+
+        return set(sorted(files))
+
     def setUp(self):
         FirefoxTestCase.setUp(self)
 
-        # Set Browser Preferences
-        self.marionette.enforce_gecko_prefs(self.browser_prefs)
+        # Force the preferences for the new profile
+        enforce_prefs = self.prefs_safebrowsing
+        enforce_prefs.update(self.prefs_provider_update_time)
+        self.marionette.enforce_gecko_prefs(enforce_prefs)
 
-        # Get safebrowsing path where downloaded data gets stored
-        self.sb_files_path = os.path.join(self.marionette.instance.profile.profile, 'safebrowsing')
+        self.safebrowsing_path = os.path.join(self.marionette.instance.profile.profile,
+                                              'safebrowsing')
+        self.safebrowsing_files = self.get_safebrowsing_files()
 
     def tearDown(self):
         try:
+            # Restart with a fresh profile
             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,
-                    ignored_exceptions=[OSError])
-
-        for data in self.test_data:
-            if self.platform not in data['platforms']:
-                continue
+        def check_downloaded(_):
+            return reduce(lambda state, pref: state and int(self.marionette.get_pref(pref)) != 1,
+                          self.prefs_provider_update_time.keys(), True)
 
-            for item in data['files']:
-                try:
-                    wait.until(
-                        lambda _: [f for f in os.listdir(self.sb_files_path)
-                                   if re.search(item, f)],
-                        message='Safe Browsing File: {} not found!'.format(item))
-                except TimeoutException:
-                    self.logger.info('Downloaded safebrowsing files: {}'.format(
-                        os.listdir(self.sb_files_path)))
-                    raise
+        try:
+            Wait(self.marionette, timeout=60).until(
+                check_downloaded, message='Not all safebrowsing files have been downloaded')
+        finally:
+            self.assertSetEqual(self.safebrowsing_files, set(os.listdir(self.safebrowsing_path)))