commit c953259f177f0e22966fab94048559c894ceda3b parent 32783c46e8ab7c6113449ee6a8793d12a02e56b7 Author: Andrea Marchesini <amarchesini@mozilla.com> Date: Thu, 23 Oct 2025 09:04:54 +0000 Bug 1994887 - Get rid of IPProtectionStates.ENROLLING, r=ip-protection-reviewers,fchasen This patch removes the "ENROLLING" state and implements a fully async enrolling logic: - when the user is signed in, we try to enroll - if not enrolled yet, when the user activates the VPN, we try to enroll again The end goal of this and next patches is to make `updateState` a sync method. Differential Revision: https://phabricator.services.mozilla.com/D268973 Diffstat:
15 files changed, 225 insertions(+), 124 deletions(-)
diff --git a/browser/components/ipprotection/IPPAutoStart.sys.mjs b/browser/components/ipprotection/IPPAutoStart.sys.mjs @@ -70,7 +70,6 @@ class IPPAutoStart { case lazy.IPProtectionStates.UNINITIALIZED: case lazy.IPProtectionStates.UNAVAILABLE: case lazy.IPProtectionStates.UNAUTHENTICATED: - case lazy.IPProtectionStates.ENROLLING: case lazy.IPProtectionStates.ERROR: this.#shouldStartWhenReady = true; break; diff --git a/browser/components/ipprotection/IPPEnrollHelper.sys.mjs b/browser/components/ipprotection/IPPEnrollHelper.sys.mjs @@ -0,0 +1,134 @@ +/* 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/. */ + +const lazy = {}; + +ChromeUtils.defineESModuleGetters(lazy, { + IPProtectionService: + "resource:///modules/ipprotection/IPProtectionService.sys.mjs", + IPPSignInWatcher: "resource:///modules/ipprotection/IPPSignInWatcher.sys.mjs", +}); + +const LOG_PREF = "browser.ipProtection.log"; + +ChromeUtils.defineLazyGetter(lazy, "logConsole", function () { + return console.createInstance({ + prefix: "IPPEnrollHelper", + maxLogLevel: Services.prefs.getBoolPref(LOG_PREF, false) ? "Debug" : "Warn", + }); +}); + +/** + * + */ +class IPPEnrollHelperSingleton { + #isEnrolled = false; + #enrolling = null; + + constructor() { + this.handleEvent = this.#handleEvent.bind(this); + } + + init() { + lazy.IPPSignInWatcher.addEventListener( + "IPPSignInWatcher:StateChanged", + this.handleEvent + ); + } + + initOnStartupCompleted() { + if (!lazy.IPPSignInWatcher.isSignedIn) { + return; + } + + try { + // This bit must be async because we want to trigger the updateState at + // the end of the rest of the initialization. + lazy.IPProtectionService.guardian + .isLinkedToGuardian(/* only cache: */ true) + .then(isEnrolled => this.#setEnrolled(isEnrolled)); + } catch (_) { + // If we are not enrolled. It's not a big deal. + } + } + + uninit() { + lazy.IPPSignInWatcher.removeEventListener( + "IPPSignInWatcher:StateChanged", + this.handleEvent + ); + } + + #handleEvent(_event) { + if (!lazy.IPPSignInWatcher.isSignedIn) { + this.#setEnrolled(false); + return; + } + + this.maybeEnroll(); + } + + async maybeEnroll() { + if (this.#isEnrolled) { + return Promise.resolve({ isEnrolled: true }); + } + + try { + this.#setEnrolled( + await lazy.IPProtectionService.guardian.isLinkedToGuardian( + /* only cache: */ false + ) + ); + } catch (error) { + return Promise.resolve({ isEnrolled: false, error: error?.message }); + } + + if (this.#isEnrolled) { + return Promise.resolve({ isEnrolled: true }); + } + + if (this.#enrolling) { + return this.#enrolling; + } + + this.#enrolling = lazy.IPProtectionService.guardian + .enroll() + .then(enrollment => { + let isEnrolled = enrollment?.ok; + + lazy.logConsole.debug( + "Guardian:", + isEnrolled ? "Enrolled" : "Enrollment Failed" + ); + + this.#setEnrolled(isEnrolled); + return { isEnrolled, error: enrollment?.error }; + }) + .catch(error => { + return { isEnrolled: false, error: error?.message }; + }) + .finally(() => { + this.#enrolling = null; + }); + + return this.#enrolling; + } + + get isEnrolled() { + return this.#isEnrolled; + } + + #setEnrolled(isEnrolled) { + if (this.#isEnrolled === isEnrolled) { + return; + } + + this.#isEnrolled = isEnrolled; + lazy.IPProtectionService.updateState(); + } +} + +const IPPEnrollHelper = new IPPEnrollHelperSingleton(); + +export { IPPEnrollHelper }; diff --git a/browser/components/ipprotection/IPPSignInWatcher.sys.mjs b/browser/components/ipprotection/IPPSignInWatcher.sys.mjs @@ -14,7 +14,7 @@ ChromeUtils.defineESModuleGetters(lazy, { * This class monitors the Sign-In state and triggers the update of the service * if needed. */ -class IPPSignInWatcherSingleton { +class IPPSignInWatcherSingleton extends EventTarget { #signedIn = false; get isSignedIn() { @@ -45,6 +45,13 @@ class IPPSignInWatcherSingleton { if (signedIn !== IPPSignInWatcher.isSignedIn) { IPPSignInWatcher.isSignedIn = signedIn; lazy.IPProtectionService.updateState(); + + IPPSignInWatcher.dispatchEvent( + new CustomEvent("IPPSignInWatcher:StateChanged", { + bubbles: true, + composed: true, + }) + ); } }, }; diff --git a/browser/components/ipprotection/IPProtection.sys.mjs b/browser/components/ipprotection/IPProtection.sys.mjs @@ -232,8 +232,6 @@ class IPProtectionWidget { * @param {Event} event - the panel shown. */ #onViewShowing(event) { - lazy.IPProtectionService.maybeEnroll(); - let { ownerGlobal } = event.target; if (this.#panels.has(ownerGlobal)) { let panel = this.#panels.get(ownerGlobal); diff --git a/browser/components/ipprotection/IPProtectionHelpers.sys.mjs b/browser/components/ipprotection/IPProtectionHelpers.sys.mjs @@ -24,6 +24,7 @@ ChromeUtils.defineESModuleGetters(lazy, { }); import { IPPAutoStartHelpers } from "resource:///modules/ipprotection/IPPAutoStart.sys.mjs"; +import { IPPEnrollHelper } from "resource:///modules/ipprotection/IPPEnrollHelper.sys.mjs"; import { IPPNimbusHelper } from "resource:///modules/ipprotection/IPPNimbusHelper.sys.mjs"; import { IPPSignInWatcher } from "resource:///modules/ipprotection/IPPSignInWatcher.sys.mjs"; import { IPPStartupCache } from "resource:///modules/ipprotection/IPPStartupCache.sys.mjs"; @@ -148,6 +149,7 @@ class VPNAddonHelper { const IPPHelpers = [ IPPStartupCache, IPPSignInWatcher, + IPPEnrollHelper, new UIHelper(), new AccountResetHelper(), new VPNAddonHelper(), diff --git a/browser/components/ipprotection/IPProtectionService.sys.mjs b/browser/components/ipprotection/IPProtectionService.sys.mjs @@ -8,6 +8,7 @@ const lazy = {}; ChromeUtils.defineESModuleGetters(lazy, { GuardianClient: "resource:///modules/ipprotection/GuardianClient.sys.mjs", + IPPEnrollHelper: "resource:///modules/ipprotection/IPPEnrollHelper.sys.mjs", IPPHelpers: "resource:///modules/ipprotection/IPProtectionHelpers.sys.mjs", IPPNimbusHelper: "resource:///modules/ipprotection/IPPNimbusHelper.sys.mjs", IPPProxyManager: "resource:///modules/ipprotection/IPPProxyManager.sys.mjs", @@ -42,9 +43,6 @@ ChromeUtils.defineLazyGetter(lazy, "logConsole", function () { * The user is not eligible (via nimbus) or still not signed in. No UI is available. * @property {string} UNAUTHENTICATED * The user is signed out but eligible (via nimbus). The panel should show the login view. - * @property {string} ENROLLING - * The user is signed in and eligible (via nimbus). The UI should show the main view, - * but not allow activation until enrollment has finished. * @property {string} READY * Ready to be activated. * @property {string} ACTIVE @@ -59,7 +57,6 @@ export const IPProtectionStates = Object.freeze({ UNINITIALIZED: "uninitialized", UNAVAILABLE: "unavailable", UNAUTHENTICATED: "unauthenticated", - ENROLLING: "enrolling", READY: "ready", ACTIVE: "active", ERROR: "error", @@ -79,7 +76,6 @@ class IPProtectionServiceSingleton extends EventTarget { #updating = false; errors = []; - enrolling = null; guardian = null; proxyManager = null; @@ -185,7 +181,6 @@ class IPProtectionServiceSingleton extends EventTarget { this.#entitlement = null; this.errors = []; - this.enrolling = null; this.#helpers.forEach(helper => helper.uninit()); @@ -206,7 +201,11 @@ class IPProtectionServiceSingleton extends EventTarget { */ async start(userAction = true) { // Wait for enrollment to finish. - await this.enrolling; + const enrollData = await lazy.IPPEnrollHelper.maybeEnroll(); + if (!enrollData || !enrollData.isEnrolled) { + this.setErrorState(enrollData.error || ERRORS.GENERIC); + return; + } // Retry getting state if the previous attempt failed. if (this.#state === IPProtectionStates.ERROR) { @@ -280,18 +279,6 @@ class IPProtectionServiceSingleton extends EventTarget { } /** - * Enroll the current account if it meets all the criteria. - * - * @returns {Promise<void>} - */ - async maybeEnroll() { - if (this.#state !== IPProtectionStates.ENROLLING) { - return null; - } - return this.#enroll(); - } - - /** * Reset the statuses that are set based on a FxA account. */ resetAccount() { @@ -305,23 +292,6 @@ class IPProtectionServiceSingleton extends EventTarget { } /** - * Checks if the user has enrolled with FxA to use the proxy. - * - * @param { boolean } onlyCached - if true only the cached clients will be checked. - * @returns {Promise<boolean>} - */ - async #isEnrolled(onlyCached) { - let isEnrolled; - try { - isEnrolled = await this.guardian.isLinkedToGuardian(onlyCached); - } catch (error) { - this.#setErrorState(error?.message); - } - - return isEnrolled; - } - - /** * Clear the current entitlement and requests a state update to dispatch * the current hasUpgraded status. * @@ -340,47 +310,6 @@ class IPProtectionServiceSingleton extends EventTarget { } /** - * Enrolls a users FxA account to use the proxy and updates the state. - * - * @returns {Promise<void>} - */ - async #enroll() { - if (this.#state !== IPProtectionStates.ENROLLING) { - return null; - } - - if (this.enrolling) { - return this.enrolling; - } - - this.enrolling = this.guardian - .enroll() - .then(enrollment => { - let ok = enrollment?.ok; - - lazy.logConsole.debug( - "Guardian:", - ok ? "Enrolled" : "Enrollment Failed" - ); - - if (!ok) { - this.#setErrorState(enrollment?.error || ERRORS.GENERIC); - return null; - } - - return this.#updateState(); - }) - .catch(error => { - this.#setErrorState(error?.message); - }) - .finally(() => { - this.enrolling = null; - }); - - return this.enrolling; - } - - /** * Gets the entitlement information for the user. */ async #getEntitlement() { @@ -460,19 +389,15 @@ class IPProtectionServiceSingleton extends EventTarget { return IPProtectionStates.READY; } - // The following are remote authentication checks and should be avoided - // whenever possible. - - // Check if the current account is enrolled with Guardian. - let enrolled = await this.#isEnrolled( - this.#state !== IPProtectionStates.ENROLLING /*onlyCached*/ - ); - if (!enrolled) { + if (!lazy.IPPEnrollHelper.isEnrolled) { return !eligible ? IPProtectionStates.UNAVAILABLE - : IPProtectionStates.ENROLLING; + : IPProtectionStates.READY; } + // The following are remote authentication checks and should be avoided + // whenever possible. + // Check if the current account can get an entitlement. let entitled = await this.#getEntitlement(); if (!entitled && !eligible) { diff --git a/browser/components/ipprotection/docs/Components.rst b/browser/components/ipprotection/docs/Components.rst @@ -32,9 +32,10 @@ A diagram of all the main components is the following: UIHelper["UI Helper"] AccountResetHelper["Account Reset Helper"] VPNAddonHelper["VPN Add-on Helper"] - IPPNimbusHelper["Nimbus Eligibility Helper"] + IPPNimbusHelper["Nimbus Eligibility Helper"] IPPAutoStart["Auto-Start Helper"] IPPEarlyStartupFilter["Early Startup Filter Helper"] + IPPEnrollHelper["Enrollment Helper"] end %% Proxy stack @@ -136,6 +137,10 @@ IPPNimbusHelper Monitors the Nimbus feature (``NimbusFeatures.ipProtection``) and triggers a state recomputation on updates. +IPPEnrollHelper + Orchestrates the user enrollment flow with Guardian and updates the service + when enrollment status changes. + How to implement new components ------------------------------- diff --git a/browser/components/ipprotection/docs/StateMachine.rst b/browser/components/ipprotection/docs/StateMachine.rst @@ -12,7 +12,6 @@ The service transitions across the following states: - ``UNINITIALIZED``: Service not initialized or feature disabled. - ``UNAVAILABLE``: User not eligible (Nimbus) or signed out with no eligibility; UI hidden. - ``UNAUTHENTICATED``: User signed out but eligible; UI shows login. -- ``ENROLLING``: User signed in and eligible; enrollment in progress. - ``READY``: Ready to activate the proxy. - ``ACTIVE``: Proxy is active. - ``ERROR``: An error occurred (see ``IPProtectionService.errors``). @@ -25,8 +24,8 @@ High‑level transitions - Not signed in → ``UNAVAILABLE`` if not eligible, otherwise ``UNAUTHENTICATED``. - Proxy already active → ``ACTIVE``. - If an entitlement is cached/valid → ``READY``. -- Otherwise, check enrollment with Guardian: - - Not enrolled → ``UNAVAILABLE`` (not eligible) or ``ENROLLING`` (eligible). +- Otherwise, check enrollment with Guardian (via ``IPPErollHelper``): + - Not enrolled → ``UNAVAILABLE`` (not eligible). - Enrolled → fetch entitlement; if successful → ``READY``, else ``UNAVAILABLE`` when not eligible. Events and integration points @@ -35,5 +34,5 @@ Events and integration points - ``IPProtectionService:StateChanged`` is dispatched on state changes with ``detail.state`` and ``detail.prevState``. - Helpers can call ``IPProtectionService.updateState()`` to request a recomputation. -- Public actions: ``start(userAction)``, ``stop(userAction)``, ``maybeEnroll()``, +- Public actions: ``start(userAction)``, ``stop(userAction)``, ``resetAccount()``, and ``startLoginFlow(browser)``. diff --git a/browser/components/ipprotection/moz.build b/browser/components/ipprotection/moz.build @@ -13,6 +13,7 @@ EXTRA_JS_MODULES.ipprotection += [ "GuardianClient.sys.mjs", "IPPAutoStart.sys.mjs", "IPPChannelFilter.sys.mjs", + "IPPEnrollHelper.sys.mjs", "IPPExceptionsManager.sys.mjs", "IPPNetworkErrorObserver.sys.mjs", "IPPNimbusHelper.sys.mjs", diff --git a/browser/components/ipprotection/tests/browser/browser_IPProtectionService.js b/browser/components/ipprotection/tests/browser/browser_IPProtectionService.js @@ -106,7 +106,7 @@ add_task(async function test_IPProtectionService_enroll() { await IPProtectionService.updateState(); Assert.equal( IPProtectionService.state, - IPProtectionStates.ENROLLING, + IPProtectionStates.READY, "User should now be enrolling" ); @@ -519,7 +519,7 @@ add_task(async function test_IPProtectionService_addon() { }, }, }); - await IPProtectionService.updateState(); + await IPProtectionService.refetchEntitlement(); const extension = ExtensionTestUtils.loadExtension({ useAddonManager: "permanent", diff --git a/browser/components/ipprotection/tests/browser/head.js b/browser/components/ipprotection/tests/browser/head.js @@ -17,6 +17,10 @@ const { IPPSignInWatcher } = ChromeUtils.importESModule( "resource:///modules/ipprotection/IPPSignInWatcher.sys.mjs" ); +const { IPPEnrollHelper } = ChromeUtils.importESModule( + "resource:///modules/ipprotection/IPPEnrollHelper.sys.mjs" +); + const { HttpServer, HTTP_403 } = ChromeUtils.importESModule( "resource://testing-common/httpd.sys.mjs" ); @@ -239,7 +243,7 @@ let DEFAULT_SERVICE_STATUS = { /* exported DEFAULT_SERVICE_STATUS */ let STUBS = { - isLinkedToGuardian: undefined, + isEnrolled: undefined, enroll: undefined, fetchUserInfo: undefined, fetchProxyPass: undefined, @@ -271,10 +275,7 @@ add_setup(async function setupVPN() { function setupStubs(stubs = STUBS) { stubs.isSignedIn = setupSandbox.stub(IPPSignInWatcher, "isSignedIn"); - stubs.isLinkedToGuardian = setupSandbox.stub( - IPProtectionService.guardian, - "isLinkedToGuardian" - ); + stubs.isEnrolled = setupSandbox.stub(IPPEnrollHelper, "isEnrolled"); stubs.enroll = setupSandbox.stub(IPProtectionService.guardian, "enroll"); stubs.fetchUserInfo = setupSandbox.stub( IPProtectionService.guardian, @@ -302,7 +303,7 @@ function setupService( } if (typeof isEnrolled != "undefined") { - stubs.isLinkedToGuardian.resolves(isEnrolled); + stubs.isEnrolled.get(() => isEnrolled); } if (typeof canEnroll != "undefined") { diff --git a/browser/components/ipprotection/tests/xpcshell/head.js b/browser/components/ipprotection/tests/xpcshell/head.js @@ -15,11 +15,13 @@ const { sinon } = ChromeUtils.importESModule( "resource://testing-common/Sinon.sys.mjs" ); -function waitForEvent(target, eventName) { +function waitForEvent(target, eventName, callback = () => true) { return new Promise(resolve => { let listener = event => { - target.removeEventListener(eventName, listener); - resolve(event); + if (callback()) { + target.removeEventListener(eventName, listener); + resolve(event); + } }; target.addEventListener(eventName, listener); }); diff --git a/browser/components/ipprotection/tests/xpcshell/test_IPProtectionPanel.js b/browser/components/ipprotection/tests/xpcshell/test_IPProtectionPanel.js @@ -12,6 +12,9 @@ const { IPProtectionService, IPProtectionStates } = ChromeUtils.importESModule( const { IPPSignInWatcher } = ChromeUtils.importESModule( "resource:///modules/ipprotection/IPPSignInWatcher.sys.mjs" ); +const { IPPEnrollHelper } = ChromeUtils.importESModule( + "resource:///modules/ipprotection/IPPEnrollHelper.sys.mjs" +); /** * A class that mocks the IP Protection panel. @@ -132,6 +135,7 @@ add_task(async function test_updateState() { add_task(async function test_IPProtectionPanel_signedIn() { let sandbox = sinon.createSandbox(); sandbox.stub(IPPSignInWatcher, "isSignedIn").get(() => true); + sandbox.stub(IPPEnrollHelper, "isEnrolled").get(() => true); sandbox .stub(IPProtectionService.guardian, "isLinkedToGuardian") .resolves(true); @@ -153,7 +157,6 @@ add_task(async function test_IPProtectionPanel_signedIn() { let signedInEventPromise = waitForEvent( IPProtectionService, "IPProtectionService:StateChanged", - false, () => IPProtectionService.state === IPProtectionStates.READY ); await IPProtectionService.updateState(); @@ -191,7 +194,6 @@ add_task(async function test_IPProtectionPanel_signedOut() { let signedOutEventPromise = waitForEvent( IPProtectionService, "IPProtectionService:StateChanged", - false, () => IPProtectionService.state === IPProtectionStates.UNAVAILABLE ); await IPProtectionService.updateState(); @@ -224,6 +226,7 @@ add_task(async function test_IPProtectionPanel_started_stopped() { let sandbox = sinon.createSandbox(); sandbox.stub(IPPSignInWatcher, "isSignedIn").get(() => true); + sandbox.stub(IPPEnrollHelper, "isEnrolled").get(() => true); sandbox .stub(IPProtectionService.guardian, "isLinkedToGuardian") .resolves(true); @@ -250,7 +253,6 @@ add_task(async function test_IPProtectionPanel_started_stopped() { let startedEventPromise = waitForEvent( IPProtectionService, "IPProtectionService:StateChanged", - false, () => IPProtectionService.state === IPProtectionStates.ACTIVE ); @@ -273,7 +275,6 @@ add_task(async function test_IPProtectionPanel_started_stopped() { let stoppedEventPromise = waitForEvent( IPProtectionService, "IPProtectionService:StateChanged", - false, () => IPProtectionService.state !== IPProtectionStates.ACTIVE ); diff --git a/browser/components/ipprotection/tests/xpcshell/test_IPProtectionService.js b/browser/components/ipprotection/tests/xpcshell/test_IPProtectionService.js @@ -15,6 +15,9 @@ const { IPProtectionService, IPProtectionStates } = ChromeUtils.importESModule( const { IPPSignInWatcher } = ChromeUtils.importESModule( "resource:///modules/ipprotection/IPPSignInWatcher.sys.mjs" ); +const { IPPEnrollHelper } = ChromeUtils.importESModule( + "resource:///modules/ipprotection/IPPEnrollHelper.sys.mjs" +); do_get_profile(); @@ -66,7 +69,6 @@ add_setup(async function () { registerCleanupFunction(async () => { await IPProtectionService.init(); - await IPProtectionService.init(); }); }); @@ -77,7 +79,13 @@ add_task(async function test_IPProtectionService_start() { let sandbox = sinon.createSandbox(); setupStubs(sandbox); - await IPProtectionService.init(); + IPProtectionService.init(); + + await waitForEvent( + IPProtectionService, + "IPProtectionService:StateChanged", + () => IPProtectionService.state === IPProtectionStates.READY + ); Assert.ok( !IPProtectionService.activatedAt, @@ -87,7 +95,6 @@ add_task(async function test_IPProtectionService_start() { let startedEventPromise = waitForEvent( IPProtectionService, "IPProtectionService:StateChanged", - false, () => IPProtectionService.state === IPProtectionStates.ACTIVE ); @@ -120,14 +127,19 @@ add_task(async function test_IPProtectionService_stop() { let sandbox = sinon.createSandbox(); setupStubs(sandbox); - await IPProtectionService.init(); + IPProtectionService.init(); + + await waitForEvent( + IPProtectionService, + "IPProtectionService:StateChanged", + () => IPProtectionService.state === IPProtectionStates.READY + ); await IPProtectionService.start(); let stoppedEventPromise = waitForEvent( IPProtectionService, "IPProtectionService:StateChanged", - false, () => IPProtectionService.state !== IPProtectionStates.ACTIVE ); IPProtectionService.stop(); @@ -156,6 +168,7 @@ add_task(async function test_IPProtectionService_stop() { */ add_task(async function test_IPProtectionService_updateState_signedIn() { let sandbox = sinon.createSandbox(); + sandbox.stub(IPPEnrollHelper, "isEnrolled").get(() => true); await IPProtectionService.init(); @@ -164,7 +177,6 @@ add_task(async function test_IPProtectionService_updateState_signedIn() { let signedInEventPromise = waitForEvent( IPProtectionService, "IPProtectionService:StateChanged", - false, () => IPProtectionService.state === IPProtectionStates.READY ); @@ -192,7 +204,6 @@ add_task(async function test_IPProtectionService_updateState_signedOut() { let signedOutEventPromise = waitForEvent( IPProtectionService, "IPProtectionService:StateChanged", - false, () => IPProtectionService.state === IPProtectionStates.UNAVAILABLE ); @@ -218,7 +229,13 @@ add_task( const sandbox = sinon.createSandbox(); setupStubs(sandbox); - await IPProtectionService.init(); + IPProtectionService.init(); + + await waitForEvent( + IPProtectionService, + "IPProtectionService:StateChanged", + () => IPProtectionService.state === IPProtectionStates.READY + ); IPProtectionService.guardian.fetchUserInfo.resolves({ status: 200, @@ -233,7 +250,6 @@ add_task( let hasUpgradedEventPromise = waitForEvent( IPProtectionService, "IPProtectionService:StateChanged", - false, () => IPProtectionService.hasUpgraded ); diff --git a/browser/components/ipprotection/tests/xpcshell/test_IPProtectionStates.js b/browser/components/ipprotection/tests/xpcshell/test_IPProtectionStates.js @@ -12,6 +12,9 @@ const { IPPNimbusHelper } = ChromeUtils.importESModule( const { IPPSignInWatcher } = ChromeUtils.importESModule( "resource:///modules/ipprotection/IPPSignInWatcher.sys.mjs" ); +const { IPPEnrollHelper } = ChromeUtils.importESModule( + "resource:///modules/ipprotection/IPPEnrollHelper.sys.mjs" +); do_get_profile(); @@ -92,6 +95,7 @@ add_task(async function test_IPProtectionStates_unauthenticated() { sandbox .stub(IPProtectionService.guardian, "isLinkedToGuardian") .resolves(false); + sandbox.stub(IPProtectionService.guardian, "enroll").resolves({ ok: true }); await IPProtectionService.init(); @@ -107,7 +111,7 @@ add_task(async function test_IPProtectionStates_unauthenticated() { Assert.equal( IPProtectionService.state, - IPProtectionStates.ENROLLING, + IPProtectionStates.READY, "IP Protection service should no longer be unauthenticated" ); @@ -146,13 +150,14 @@ add_task(async function test_IPProtectionStates_enrolling() { Assert.equal( IPProtectionService.state, - IPProtectionStates.ENROLLING, - "IP Protection service should be enrolling" + IPProtectionStates.READY, + "IP Protection service should be ready" ); IPProtectionService.guardian.isLinkedToGuardian.resolves(true); - await IPProtectionService.maybeEnroll(); + const enrollData = await IPPEnrollHelper.maybeEnroll(); + Assert.ok(enrollData.isEnrolled, "Fully enrolled"); Assert.equal( IPProtectionService.state, @@ -224,7 +229,13 @@ add_task(async function test_IPProtectionStates_active() { }, }); - await IPProtectionService.init(); + IPProtectionService.init(); + + await waitForEvent( + IPProtectionService, + "IPProtectionService:StateChanged", + () => IPProtectionService.state === IPProtectionStates.READY + ); Assert.equal( IPProtectionService.state,