Commit 6502db91 authored by Beth Rennie's avatar Beth Rennie Committed by brennie@mozilla.com
Browse files

Bug 1956080 - Add NimbusTestUtils helpers for dealing with the SQL database... 

Bug 1956080 - Add NimbusTestUtils helpers for dealing with the SQL database r=nimbus-reviewers,relud

When we start executing SQL queries during unenrollment, the
ExperimentManager.unenroll() function will become effectively async. We
need to change our cleanup functions to work in this semi-async world.

Differential Revision: https://phabricator.services.mozilla.com/D250504
parent 5fabb062

toolkit/components/nimbus/test/NimbusTestUtils.sys.mjs

+45 −7
Changes for toolkit/components/nimbus/test/NimbusTestUtils.sys.mjs: 45 added lines, 7 removed lines.
Original line number Diff line number Diff line
@@ -16,6 +16,8 @@ ChromeUtils.defineESModuleGetters(lazy, { 
  JsonSchema: "resource://gre/modules/JsonSchema.sys.mjs",

  NetUtil: "resource://gre/modules/NetUtil.sys.mjs",

  ExperimentManager: "resource://nimbus/lib/ExperimentManager.sys.mjs",

  ProfilesDatastoreService:

    "moz-src:///toolkit/profile/ProfilesDatastoreService.sys.mjs",

  RemoteSettingsExperimentLoader:

    "resource://nimbus/lib/RemoteSettingsExperimentLoader.sys.mjs",

  sinon: "resource://testing-common/Sinon.sys.mjs",
@@ -156,6 +158,12 @@ export const NimbusTestUtils = { 
      );



      NimbusTestUtils.cleanupStorePrefCache();



      // TODO(bug 1956082): This is an async method that we are not awaiting.

      //

      // Only browser tests and tests that otherwise have manually enabled the

      // ProfilesDatastoreService need to await the result.

      return NimbusTestUtils.cleanupEnrollmentDatabase();

    },

  },
@@ -371,14 +379,37 @@ export const NimbusTestUtils = { 
   *         The ExperimentManager to clean up. Defaults to the global

   *         ExperimentManager.

   */

  cleanupManager(slugs, { manager } = {}) {

  async cleanupManager(slugs, { manager } = {}) {

    const experimentManager = manager ?? ExperimentAPI.manager;



    for (const slug of slugs) {

      experimentManager.unenroll(slug);

      await experimentManager.unenroll(slug);

    }



    NimbusTestUtils.assert.storeIsEmpty(experimentManager.store);

    await NimbusTestUtils.assert.storeIsEmpty(experimentManager.store);

  },



  async cleanupEnrollmentDatabase() {

    const conn = await lazy.ProfilesDatastoreService.getConnection();

    if (!conn) {

      // We are in an xpcshell test that has not initialized the

      // ProfilesDatastoreService.

      //

      // TODO(bug 1967779): require the ProfilesDatastoreService to be initialized

      // and remove this check.

      return;

    }



    // TODO(bug 1956082): This should filter only active = false enrollments,

    // but that requires the unenrollment flow to be fully async.

    await conn.execute(

      `

        DELETE FROM NimbusEnrollments

        WHERE

          profileId = :profileId;

      `,

      { profileId: ExperimentAPI.profileId }

    );

  },



  /**
@@ -443,8 +474,14 @@ export const NimbusTestUtils = { 
    experimentManager.store._syncToChildren({ flush: true });



    return function doEnrollmentCleanup() {

      experimentManager.unenroll(enrollment.slug);

      // TODO(bug 1956082): This is an async method that we are not awaiting.

      //

      // Only browser tests and tests that otherwise have manually enabled the

      // ProfilesDatastoreService need to await the result.

      const promise = experimentManager.unenroll(enrollment.slug);

      experimentManager.store._deleteForTests(enrollment.slug);



      return promise;

    };

  },
@@ -527,7 +564,7 @@ export const NimbusTestUtils = { 
   *         The store to delete.

   */

  async removeStore(store) {

    NimbusTestUtils.assert.storeIsEmpty(store);

    await NimbusTestUtils.assert.storeIsEmpty(store);



    // Prevent the next save from happening.

    store._store._saver.disarm();
@@ -646,8 +683,9 @@ export const NimbusTestUtils = { 
      sandbox,

      loader,

      manager,

      cleanup() {

        NimbusTestUtils.assert.storeIsEmpty(manager.store);

      async cleanup() {

        await NimbusTestUtils.assert.storeIsEmpty(manager.store);



        ExperimentAPI._resetForTests();

        sandbox.restore();