Bug 1377104 - Should clear all stored site data dynamically, r?Gijs draft
authorFischer.json <fischer.json@gmail.com>
Tue, 22 Aug 2017 14:57:43 +0800
changeset 651866 d967dea7cad51a4f39622b7c0c86bb1458ff3aeb
parent 650228 128a79130ecd6f277190d031a623f991c73c5272
child 727888 75e763cec7282e3a054b5ff26631f5369b83ec16
push id75836
push userbmo:fliu@mozilla.com
push dateThu, 24 Aug 2017 04:21:25 +0000
reviewersGijs
bugs1377104
milestone57.0a1
Bug 1377104 - Should clear all stored site data dynamically, r?Gijs MozReview-Commit-ID: 69dWoxw2gGm
browser/components/preferences/SiteDataManager.jsm
browser/components/preferences/in-content-new/tests/browser_siteData2.js
browser/components/preferences/in-content-new/tests/head.js
browser/components/preferences/in-content/tests/browser_siteData2.js
browser/components/preferences/in-content/tests/head.js
--- a/browser/components/preferences/SiteDataManager.jsm
+++ b/browser/components/preferences/SiteDataManager.jsm
@@ -28,69 +28,64 @@ this.SiteDataManager = {
   //   - quotaUsage: the usage of indexedDB and localStorage.
   //   - appCacheList: an array of app cache; instances of nsIApplicationCache
   _sites: new Map(),
 
   _getQuotaUsagePromise: null,
 
   _quotaUsageRequest: null,
 
-  updateSites() {
+  async updateSites() {
     Services.obs.notifyObservers(null, "sitedatamanager:updating-sites");
+    await this._getQuotaUsage();
+    this._updateAppCache();
+    Services.obs.notifyObservers(null, "sitedatamanager:sites-updated");
+  },
 
+  _getQuotaUsage() {
     // Clear old data and requests first
     this._sites.clear();
     this._cancelGetQuotaUsage();
-
-    this._getQuotaUsage()
-        .then(results => {
-          for (let result of results) {
-            let principal =
-              Services.scriptSecurityManager.createCodebasePrincipalFromOrigin(result.origin);
-            let uri = principal.URI;
-            if (uri.scheme == "http" || uri.scheme == "https") {
-              let site = this._sites.get(uri.host);
-              if (!site) {
-                site = {
-                  persisted: false,
-                  quotaUsage: 0,
-                  principals: [],
-                  appCacheList: [],
-                };
-              }
-              // Assume 3 sites:
-              //   - Site A (not persisted): https://www.foo.com
-              //   - Site B (not persisted): https://www.foo.com^userContextId=2
-              //   - Site C (persisted):     https://www.foo.com:1234
-              // Although only C is persisted, grouping by host, as a result,
-              // we still mark as persisted here under this host group.
-              if (result.persisted) {
-                site.persisted = true;
-              }
-              site.principals.push(principal);
-              site.quotaUsage += result.usage;
-              this._sites.set(uri.host, site);
+    this._getQuotaUsagePromise = new Promise(resolve => {
+      let onUsageResult = request => {
+        let items = request.result;
+        for (let item of items) {
+          let principal =
+            Services.scriptSecurityManager.createCodebasePrincipalFromOrigin(item.origin);
+          let uri = principal.URI;
+          if (uri.scheme == "http" || uri.scheme == "https") {
+            let site = this._sites.get(uri.host);
+            if (!site) {
+              site = {
+                persisted: false,
+                quotaUsage: 0,
+                principals: [],
+                appCacheList: [],
+              };
             }
+            // Assume 3 sites:
+            //   - Site A (not persisted): https://www.foo.com
+            //   - Site B (not persisted): https://www.foo.com^userContextId=2
+            //   - Site C (persisted):     https://www.foo.com:1234
+            // Although only C is persisted, grouping by host, as a result,
+            // we still mark as persisted here under this host group.
+            if (item.persisted) {
+              site.persisted = true;
+            }
+            site.principals.push(principal);
+            site.quotaUsage += item.usage;
+            this._sites.set(uri.host, site);
           }
-          this._updateAppCache();
-          Services.obs.notifyObservers(null, "sitedatamanager:sites-updated");
-        });
-  },
-
-  _getQuotaUsage() {
-    this._getQuotaUsagePromise = new Promise(resolve => {
-      let callback = {
-        onUsageResult(request) {
-          resolve(request.result);
         }
+        resolve();
       };
       // XXX: The work of integrating localStorage into Quota Manager is in progress.
       //      After the bug 742822 and 1286798 landed, localStorage usage will be included.
       //      So currently only get indexedDB usage.
-      this._quotaUsageRequest = this._qms.getUsage(callback);
+      this._quotaUsageRequest = this._qms.getUsage(onUsageResult);
     });
     return this._getQuotaUsagePromise;
   },
 
   _cancelGetQuotaUsage() {
     if (this._quotaUsageRequest) {
       this._quotaUsageRequest.cancel();
       this._quotaUsageRequest = null;
@@ -232,26 +227,37 @@ this.SiteDataManager = {
     if (promises.length > 0) {
       Promise.all(promises).then(() => this.updateSites());
     }
     if (unknownHost) {
       throw `SiteDataManager: removing unknown site of ${unknownHost}`;
     }
   },
 
-  removeAll() {
+  async removeAll() {
+    Services.cache2.clear();
+    Services.cookies.removeAll();
+    OfflineAppCacheHelper.clear();
+    // Refresh sites using quota usage again.
+    // This is for the case:
+    //   1. User goes to the about:preferences Site Data section.
+    //   2. With the about:preferences opened, user visits another website.
+    //   3. The website saves to quota usage, like indexedDB.
+    //   4. User goes back to the Site Data section and commands to clear all site data.
+    // For this case, we should refresh the site list so not to miss the website in the step 3.
+    // We don't do "Clear All" on the quota manager like the cookie, appcache, http cache above
+    // because that would clear browser data as well too,
+    // see https://bugzilla.mozilla.org/show_bug.cgi?id=1312361#c9
+    await this._getQuotaUsage();
     let promises = [];
     for (let site of this._sites.values()) {
       this._removePermission(site);
       promises.push(this._removeQuotaUsage(site));
     }
-    Services.cache2.clear();
-    Services.cookies.removeAll();
-    OfflineAppCacheHelper.clear();
-    Promise.all(promises).then(() => this.updateSites());
+    return Promise.all(promises).then(() => this.updateSites());
   },
 
   isPrivateCookie(cookie) {
     let { userContextId } = cookie.originAttributes;
     // A private cookie is when its userContextId points to a private identity.
     return userContextId && !ContextualIdentityService.getPublicIdentityFromId(userContextId);
   }
 };
--- a/browser/components/preferences/in-content-new/tests/browser_siteData2.js
+++ b/browser/components/preferences/in-content-new/tests/browser_siteData2.js
@@ -13,16 +13,27 @@ function promiseSettingsDialogClose() {
       if (dialogWin.document.documentURI === "chrome://browser/content/preferences/siteDataSettings.xul") {
         isnot(dialogOverlay.style.visibility, "visible", "The Settings dialog should be hidden");
         resolve();
       }
     }, { once: true });
   });
 }
 
+function assertAllSitesNotListed(win) {
+  let frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
+  let removeBtn = frameDoc.getElementById("removeSelected");
+  let removeAllBtn = frameDoc.getElementById("removeAll");
+  let sitesList = frameDoc.getElementById("sitesList");
+  let sites = sitesList.getElementsByTagName("richlistitem");
+  is(sites.length, 0, "Should not list all sites");
+  is(removeBtn.disabled, true, "Should disable the removeSelected button");
+  is(removeAllBtn.disabled, true, "Should disable the removeAllBtn button");
+}
+
 // Test selecting and removing all sites one by one
 add_task(async function() {
   await SpecialPowers.pushPrefEnv({set: [["browser.storageManager.enabled", true]]});
   mockSiteDataManager.register(SiteDataManager);
   mockSiteDataManager.fakeSites = [
     {
       usage: 1024,
       principal: Services.scriptSecurityManager
@@ -65,74 +76,63 @@ add_task(async function() {
   // Test the initial state
   assertSitesListed(doc, fakeHosts);
 
   // Test the "Cancel" button
   settingsDialogClosePromise = promiseSettingsDialogClose();
   frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
   cancelBtn = frameDoc.getElementById("cancel");
   removeAllSitesOneByOne();
-  assertAllSitesNotListed();
+  assertAllSitesNotListed(win);
   cancelBtn.doCommand();
   await settingsDialogClosePromise;
   await openSiteDataSettingsDialog();
   assertSitesListed(doc, fakeHosts);
 
   // Test the "Save Changes" button but cancelling save
   let cancelPromise = promiseAlertDialogOpen("cancel");
   settingsDialogClosePromise = promiseSettingsDialogClose();
   frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
   saveBtn = frameDoc.getElementById("save");
   removeAllSitesOneByOne();
-  assertAllSitesNotListed();
+  assertAllSitesNotListed(win);
   saveBtn.doCommand();
   await cancelPromise;
   await settingsDialogClosePromise;
   await openSiteDataSettingsDialog();
   assertSitesListed(doc, fakeHosts);
 
   // Test the "Save Changes" button and accepting save
   let acceptPromise = promiseAlertDialogOpen("accept");
   settingsDialogClosePromise = promiseSettingsDialogClose();
   updatePromise = promiseSiteDataManagerSitesUpdated();
   frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
   saveBtn = frameDoc.getElementById("save");
   removeAllSitesOneByOne();
-  assertAllSitesNotListed();
+  assertAllSitesNotListed(win);
   saveBtn.doCommand();
   await acceptPromise;
   await settingsDialogClosePromise;
   await updatePromise;
   await openSiteDataSettingsDialog();
-  assertAllSitesNotListed();
+  assertAllSitesNotListed(win);
 
   mockSiteDataManager.unregister();
   await BrowserTestUtils.removeTab(gBrowser.selectedTab);
 
   function removeAllSitesOneByOne() {
     frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
     let removeBtn = frameDoc.getElementById("removeSelected");
     let sitesList = frameDoc.getElementById("sitesList");
     let sites = sitesList.getElementsByTagName("richlistitem");
     for (let i = sites.length - 1; i >= 0; --i) {
       sites[i].click();
       removeBtn.doCommand();
     }
   }
-
-  function assertAllSitesNotListed() {
-    frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
-    let removeBtn = frameDoc.getElementById("removeSelected");
-    let removeAllBtn = frameDoc.getElementById("removeAll");
-    let sitesList = frameDoc.getElementById("sitesList");
-    let sites = sitesList.getElementsByTagName("richlistitem");
-    is(sites.length, 0, "Should not list all sites");
-    is(removeBtn.disabled, true, "Should disable the removeSelected button");
-    is(removeAllBtn.disabled, true, "Should disable the removeAllBtn button");
-  }
 });
 
 // Test selecting and removing partial sites
 add_task(async function() {
   await SpecialPowers.pushPrefEnv({set: [["browser.storageManager.enabled", true]]});
   mockSiteDataManager.register(SiteDataManager);
   mockSiteDataManager.fakeSites = [
     {
@@ -371,8 +371,69 @@ add_task(async function() {
 
   expected = prefStrBundle.getString("persistent");
   let status = siteItems[0].getAttribute("status");
   is(status, expected, "Should mark persisted status across scheme, port and origin attributes");
 
   mockSiteDataManager.unregister();
   await BrowserTestUtils.removeTab(gBrowser.selectedTab);
 });
+
+// Test dynamically clearing all site data
+add_task(async function() {
+  await SpecialPowers.pushPrefEnv({set: [["browser.storageManager.enabled", true]]});
+  mockSiteDataManager.register(SiteDataManager);
+  mockSiteDataManager.fakeSites = [
+    {
+      usage: 1024,
+      principal: Services.scriptSecurityManager
+                         .createCodebasePrincipalFromOrigin("https://account.xyz.com"),
+      persisted: true
+    },
+    {
+      usage: 1024,
+      principal: Services.scriptSecurityManager
+                         .createCodebasePrincipalFromOrigin("https://shopping.xyz.com"),
+      persisted: false
+    },
+  ];
+  let fakeHosts = mockSiteDataManager.fakeSites.map(site => site.principal.URI.host);
+
+  // Test the initial state
+  let updatePromise = promiseSiteDataManagerSitesUpdated();
+  await openPreferencesViaOpenPreferencesAPI("privacy", { leaveOpen: true });
+  await updatePromise;
+  await openSiteDataSettingsDialog();
+  let doc = gBrowser.selectedBrowser.contentDocument;
+  assertSitesListed(doc, fakeHosts);
+
+  // Add more sites dynamically
+  mockSiteDataManager.fakeSites.push({
+    usage: 1024,
+    principal: Services.scriptSecurityManager
+                       .createCodebasePrincipalFromOrigin("http://cinema.bar.com"),
+    persisted: true
+  }, {
+    usage: 1024,
+    principal: Services.scriptSecurityManager
+                       .createCodebasePrincipalFromOrigin("http://email.bar.com"),
+    persisted: false
+  });
+
+  // Test clearing all site data dynamically
+  let win = gBrowser.selectedBrowser.contentWindow;
+  let frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
+  updatePromise = promiseSiteDataManagerSitesUpdated();
+  let acceptRemovePromise = promiseAlertDialogOpen("accept");
+  let settingsDialogClosePromise = promiseSettingsDialogClose();
+  let removeAllBtn = frameDoc.getElementById("removeAll");
+  let saveBtn = frameDoc.getElementById("save");
+  removeAllBtn.doCommand();
+  saveBtn.doCommand();
+  await acceptRemovePromise;
+  await settingsDialogClosePromise;
+  await updatePromise;
+  await openSiteDataSettingsDialog();
+  assertAllSitesNotListed(win);
+
+  mockSiteDataManager.unregister();
+  await BrowserTestUtils.removeTab(gBrowser.selectedTab);
+});
--- a/browser/components/preferences/in-content-new/tests/head.js
+++ b/browser/components/preferences/in-content-new/tests/head.js
@@ -247,45 +247,41 @@ async function evaluateSearchResults(key
       is_element_hidden(child, "Should not be in search results");
     }
   }
 }
 
 const mockSiteDataManager = {
 
   _SiteDataManager: null,
-  _originalGetQuotaUsage: null,
+  _originalQMS: null,
   _originalRemoveQuotaUsage: null,
 
-  _getQuotaUsage() {
-    let results = [];
-    this.fakeSites.forEach(site => {
-      results.push({
-        origin: site.principal.origin,
-        usage: site.usage,
-        persisted: site.persisted
-      });
-    });
-    this._SiteDataManager._getQuotaUsagePromise = Promise.resolve(results);
-    return this._SiteDataManager._getQuotaUsagePromise;
+  getUsage(onUsageResult) {
+    let result = this.fakeSites.map(site => ({
+      origin: site.principal.origin,
+      usage: site.usage,
+      persisted: site.persisted
+    }));
+    onUsageResult({ result });
   },
 
   _removeQuotaUsage(site) {
     var target = site.principals[0].URI.host;
     this.fakeSites = this.fakeSites.filter(fakeSite => {
       return fakeSite.principal.URI.host != target;
     });
   },
 
   register(SiteDataManager) {
     this._SiteDataManager = SiteDataManager;
-    this._originalGetQuotaUsage = this._SiteDataManager._getQuotaUsage;
-    this._SiteDataManager._getQuotaUsage = this._getQuotaUsage.bind(this);
+    this._originalQMS = this._SiteDataManager._qms;
+    this._SiteDataManager._qms = this;
     this._originalRemoveQuotaUsage = this._SiteDataManager._removeQuotaUsage;
     this._SiteDataManager._removeQuotaUsage = this._removeQuotaUsage.bind(this);
     this.fakeSites = null;
   },
 
   unregister() {
-    this._SiteDataManager._getQuotaUsage = this._originalGetQuotaUsage;
+    this._SiteDataManager._qms = this._originalQMS;
     this._SiteDataManager._removeQuotaUsage = this._originalRemoveQuotaUsage;
   }
 };
--- a/browser/components/preferences/in-content/tests/browser_siteData2.js
+++ b/browser/components/preferences/in-content/tests/browser_siteData2.js
@@ -1,10 +1,22 @@
 "use strict";
 
+
+function assertAllSitesNotListed(win) {
+  let frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
+  let removeBtn = frameDoc.getElementById("removeSelected");
+  let removeAllBtn = frameDoc.getElementById("removeAll");
+  let sitesList = frameDoc.getElementById("sitesList");
+  let sites = sitesList.getElementsByTagName("richlistitem");
+  is(sites.length, 0, "Should not list all sites");
+  is(removeBtn.disabled, true, "Should disable the removeSelected button");
+  is(removeAllBtn.disabled, true, "Should disable the removeAllBtn button");
+}
+
 // Test selecting and removing all sites one by one
 add_task(async function() {
   await SpecialPowers.pushPrefEnv({set: [["browser.storageManager.enabled", true]]});
   mockSiteDataManager.register();
   mockSiteDataManager.fakeSites = [
     {
       usage: 1024,
       principal: Services.scriptSecurityManager
@@ -47,74 +59,63 @@ add_task(async function() {
   // Test the initial state
   assertSitesListed(doc, fakeHosts);
 
   // Test the "Cancel" button
   settingsDialogClosePromise = promiseSettingsDialogClose();
   frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
   cancelBtn = frameDoc.getElementById("cancel");
   removeAllSitesOneByOne();
-  assertAllSitesNotListed();
+  assertAllSitesNotListed(win);
   cancelBtn.doCommand();
   await settingsDialogClosePromise;
   await openSettingsDialog();
   assertSitesListed(doc, fakeHosts);
 
   // Test the "Save Changes" button but cancelling save
   let cancelPromise = promiseAlertDialogOpen("cancel");
   settingsDialogClosePromise = promiseSettingsDialogClose();
   frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
   saveBtn = frameDoc.getElementById("save");
   removeAllSitesOneByOne();
-  assertAllSitesNotListed();
+  assertAllSitesNotListed(win);
   saveBtn.doCommand();
   await cancelPromise;
   await settingsDialogClosePromise;
   await openSettingsDialog();
   assertSitesListed(doc, fakeHosts);
 
   // Test the "Save Changes" button and accepting save
   let acceptPromise = promiseAlertDialogOpen("accept");
   settingsDialogClosePromise = promiseSettingsDialogClose();
   updatePromise = promiseSitesUpdated();
   frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
   saveBtn = frameDoc.getElementById("save");
   removeAllSitesOneByOne();
-  assertAllSitesNotListed();
+  assertAllSitesNotListed(win);
   saveBtn.doCommand();
   await acceptPromise;
   await settingsDialogClosePromise;
   await updatePromise;
   await openSettingsDialog();
-  assertAllSitesNotListed();
+  assertAllSitesNotListed(win);
 
   mockSiteDataManager.unregister();
   await BrowserTestUtils.removeTab(gBrowser.selectedTab);
 
   function removeAllSitesOneByOne() {
     frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
     let removeBtn = frameDoc.getElementById("removeSelected");
     let sitesList = frameDoc.getElementById("sitesList");
     let sites = sitesList.getElementsByTagName("richlistitem");
     for (let i = sites.length - 1; i >= 0; --i) {
       sites[i].click();
       removeBtn.doCommand();
     }
   }
-
-  function assertAllSitesNotListed() {
-    frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
-    let removeBtn = frameDoc.getElementById("removeSelected");
-    let removeAllBtn = frameDoc.getElementById("removeAll");
-    let sitesList = frameDoc.getElementById("sitesList");
-    let sites = sitesList.getElementsByTagName("richlistitem");
-    is(sites.length, 0, "Should not list all sites");
-    is(removeBtn.disabled, true, "Should disable the removeSelected button");
-    is(removeAllBtn.disabled, true, "Should disable the removeAllBtn button");
-  }
 });
 
 // Test selecting and removing partial sites
 add_task(async function() {
   await SpecialPowers.pushPrefEnv({set: [["browser.storageManager.enabled", true]]});
   mockSiteDataManager.register();
   mockSiteDataManager.fakeSites = [
     {
@@ -292,8 +293,69 @@ add_task(async function() {
   await settingsDialogClosePromise;
   await updatePromise;
   await openSettingsDialog();
   assertSitesListed(doc, fakeHosts.filter(host => !host.includes("xyz")));
 
   mockSiteDataManager.unregister();
   await BrowserTestUtils.removeTab(gBrowser.selectedTab);
 });
+
+// Test dynamically clearing all site data
+add_task(async function() {
+  await SpecialPowers.pushPrefEnv({set: [["browser.storageManager.enabled", true]]});
+  mockSiteDataManager.register();
+  mockSiteDataManager.fakeSites = [
+    {
+      usage: 1024,
+      principal: Services.scriptSecurityManager
+                         .createCodebasePrincipalFromOrigin("https://account.xyz.com"),
+      persisted: true
+    },
+    {
+      usage: 1024,
+      principal: Services.scriptSecurityManager
+                         .createCodebasePrincipalFromOrigin("https://shopping.xyz.com"),
+      persisted: false
+    },
+  ];
+  let fakeHosts = mockSiteDataManager.fakeSites.map(site => site.principal.URI.host);
+
+  // Test the initial state
+  let updatePromise = promiseSitesUpdated();
+  await openPreferencesViaOpenPreferencesAPI("advanced", "networkTab", { leaveOpen: true });
+  await updatePromise;
+  await openSettingsDialog();
+  let doc = gBrowser.selectedBrowser.contentDocument;
+  assertSitesListed(doc, fakeHosts);
+
+  // Add more sites dynamically
+  mockSiteDataManager.fakeSites.push({
+    usage: 1024,
+    principal: Services.scriptSecurityManager
+                       .createCodebasePrincipalFromOrigin("http://cinema.bar.com"),
+    persisted: true
+  }, {
+    usage: 1024,
+    principal: Services.scriptSecurityManager
+                       .createCodebasePrincipalFromOrigin("http://email.bar.com"),
+    persisted: false
+  });
+
+  // Test clearing all site data dynamically
+  let win = gBrowser.selectedBrowser.contentWindow;
+  let frameDoc = win.gSubDialog._topDialog._frame.contentDocument;
+  updatePromise = promiseSitesUpdated();
+  let acceptRemovePromise = promiseAlertDialogOpen("accept");
+  let settingsDialogClosePromise = promiseSettingsDialogClose();
+  let removeAllBtn = frameDoc.getElementById("removeAll");
+  let saveBtn = frameDoc.getElementById("save");
+  removeAllBtn.doCommand();
+  saveBtn.doCommand();
+  await acceptRemovePromise;
+  await settingsDialogClosePromise;
+  await updatePromise;
+  await openSettingsDialog();
+  assertAllSitesNotListed(win);
+
+  mockSiteDataManager.unregister();
+  await BrowserTestUtils.removeTab(gBrowser.selectedTab);
+});
--- a/browser/components/preferences/in-content/tests/head.js
+++ b/browser/components/preferences/in-content/tests/head.js
@@ -12,49 +12,45 @@ registerCleanupFunction(function() {
   Services.prefs.clearUserPref("browser.preferences.useOldOrganization");
 });
 
 const kDefaultWait = 2000;
 const REMOVE_DIALOG_URL = "chrome://browser/content/preferences/siteDataRemoveSelected.xul";
 const { SiteDataManager } = Cu.import("resource:///modules/SiteDataManager.jsm", {});
 const mockSiteDataManager = {
 
-  _originalGetQuotaUsage: null,
+  _originalQMS: null,
   _originalRemoveQuotaUsage: null,
 
-  _getQuotaUsage() {
-    let results = [];
-    this.fakeSites.forEach(site => {
-      results.push({
-        origin: site.principal.origin,
-        usage: site.usage,
-        persisted: site.persisted
-      });
-    });
-    SiteDataManager._getQuotaUsagePromise = Promise.resolve(results);
-    return SiteDataManager._getQuotaUsagePromise;
+  getUsage(onUsageResult) {
+    let result = this.fakeSites.map(site => ({
+      origin: site.principal.origin,
+      usage: site.usage,
+      persisted: site.persisted
+    }));
+    onUsageResult({ result });
   },
 
   _removeQuotaUsage(site) {
     var target = site.principals[0].URI.host;
     this.fakeSites = this.fakeSites.filter(fakeSite => {
       return fakeSite.principal.URI.host != target;
     });
   },
 
   register() {
-    this._originalGetQuotaUsage = SiteDataManager._getQuotaUsage;
-    SiteDataManager._getQuotaUsage = this._getQuotaUsage.bind(this);
+    this._originalQMS = SiteDataManager._qms;
+    SiteDataManager._qms = this;
     this._originalRemoveQuotaUsage = SiteDataManager._removeQuotaUsage;
     SiteDataManager._removeQuotaUsage = this._removeQuotaUsage.bind(this);
     this.fakeSites = null;
   },
 
   unregister() {
-    SiteDataManager._getQuotaUsage = this._originalGetQuotaUsage;
+    SiteDataManager._qms = this._originalQMS;
     SiteDataManager._removeQuotaUsage = this._originalRemoveQuotaUsage;
   }
 };
 
 function is_hidden(aElement) {
   var style = aElement.ownerGlobal.getComputedStyle(aElement);
   if (style.display == "none")
     return true;