|  | 1 | +from marionette_driver import By, Wait, errors
 | 
|  | 2 | +from marionette_driver.localization import L10n
 | 
|  | 3 | +from marionette_harness import MarionetteTestCase
 | 
|  | 4 | +
 | 
|  | 5 | +NETWORK_CHECK_URL = "https://check.torproject.org/"
 | 
|  | 6 | +TOR_BOOTSTRAP_TIMEOUT = 30  # 30s
 | 
|  | 7 | +
 | 
|  | 8 | +STRINGS_LOCATION = "chrome://torbutton/locale/torConnect.properties"
 | 
|  | 9 | +
 | 
|  | 10 | +
 | 
|  | 11 | +class TestNetworkCheck(MarionetteTestCase):
 | 
|  | 12 | +    def setUp(self):
 | 
|  | 13 | +        MarionetteTestCase.setUp(self)
 | 
|  | 14 | +
 | 
|  | 15 | +        self.l10n = L10n(self.marionette)
 | 
|  | 16 | +
 | 
|  | 17 | +    def attemptConnection(self, tries=1):
 | 
|  | 18 | +        if tries > 3:
 | 
|  | 19 | +            self.assertTrue(False, "Failed to connect to Tor after 3 attempts")
 | 
|  | 20 | +
 | 
|  | 21 | +        connectBtn = self.marionette.find_element(By.ID, "connectButton")
 | 
|  | 22 | +        Wait(self.marionette, timeout=10).until(
 | 
|  | 23 | +            lambda _: connectBtn.is_displayed(),
 | 
|  | 24 | +            message="Timed out waiting for tor connect button to show up.",
 | 
|  | 25 | +        )
 | 
|  | 26 | +        connectBtn.click()
 | 
|  | 27 | +
 | 
|  | 28 | +        try:
 | 
|  | 29 | +
 | 
|  | 30 | +            def check(m):
 | 
|  | 31 | +                if not m.get_url().startswith("about:torconnect"):
 | 
|  | 32 | +                    # We have finished connecting and have been redirected.
 | 
|  | 33 | +                    return True
 | 
|  | 34 | +
 | 
|  | 35 | +                try:
 | 
|  | 36 | +                    heading = self.marionette.find_element(By.ID, "tor-connect-heading")
 | 
|  | 37 | +                except errors.NoSuchElementException:
 | 
|  | 38 | +                    # Page is probably redirecting.
 | 
|  | 39 | +                    return False
 | 
|  | 40 | +
 | 
|  | 41 | +                if heading.text not in [
 | 
|  | 42 | +                    self.l10n.localize_property(
 | 
|  | 43 | +                        [STRINGS_LOCATION], "torConnect.torConnecting"
 | 
|  | 44 | +                    ),
 | 
|  | 45 | +                    self.l10n.localize_property(
 | 
|  | 46 | +                        [STRINGS_LOCATION], "torConnect.torConnected"
 | 
|  | 47 | +                    ),
 | 
|  | 48 | +                ]:
 | 
|  | 49 | +                    raise ValueError("Tor connect page is not connecting or connected")
 | 
|  | 50 | +
 | 
|  | 51 | +                return False
 | 
|  | 52 | +
 | 
|  | 53 | +            Wait(self.marionette, timeout=TOR_BOOTSTRAP_TIMEOUT).until(check)
 | 
|  | 54 | +        except (errors.TimeoutException, ValueError):
 | 
|  | 55 | +            cancelBtn = self.marionette.find_element(By.ID, "cancelButton")
 | 
|  | 56 | +            if cancelBtn.is_displayed():
 | 
|  | 57 | +                cancelBtn.click()
 | 
|  | 58 | +
 | 
|  | 59 | +            self.attemptConnection(tries + 1)
 | 
|  | 60 | +
 | 
|  | 61 | +    def test_network_check(self):
 | 
|  | 62 | +        self.attemptConnection()
 | 
|  | 63 | +        self.marionette.navigate(NETWORK_CHECK_URL)
 | 
|  | 64 | +        self.assertRegex(
 | 
|  | 65 | +            self.marionette.title,
 | 
|  | 66 | +            r"^Congratulations\.",
 | 
|  | 67 | +            f"{NETWORK_CHECK_URL} should have the expected title.",
 | 
|  | 68 | +        ) |