Fix Bug 1511804 - CFRPageActions tries to wait for a layout flush but doesn't (#4585)

lib/CFRPageActions.jsm

@@ -68,8 +68,7 @@ class PageAction {
     // Wait for layout to flush to avoid a synchronous reflow then calculate the
     // label width. We can safely get the width even though the recommendation is
     // collapsed; the label itself remains full width (with its overflow hidden)
-    await this.window.promiseDocumentFlushed;
-    const [{width}] = this.label.getClientRects();
+    let [{width}] = await this.window.promiseDocumentFlushed(() => this.label.getClientRects());
     this.urlbar.style.setProperty("--cfr-label-width", `${width}px`);
 
     this.container.addEventListener("click", this._handleClick);

test/unit/asrouter/CFRPageActions.test.js

@@ -84,7 +84,7 @@ describe("CFRPageActions", () => {
     globals = new GlobalOverrider();
     globals.set({
       Localization: class {},
-      promiseDocumentFlushed: Promise.resolve(),
+      promiseDocumentFlushed: sandbox.stub().callsFake(fn => Promise.resolve(fn())),
       PopupNotifications: {
         show: sandbox.stub(),
         remove: sandbox.stub(),
@@ -131,11 +131,9 @@ describe("CFRPageActions", () => {
       });
       it("should wait for the document layout to flush", async () => {
         sandbox.spy(pageAction.label, "getClientRects");
-        const promiseGetterStub = sandbox.stub().resolves();
-        sandbox.stub(global, "promiseDocumentFlushed").get(promiseGetterStub);
         await pageAction.show(fakeRecommendation);
-        assert.calledOnce(promiseGetterStub);
-        assert.callOrder(promiseGetterStub, pageAction.label.getClientRects);
+        assert.calledOnce(global.promiseDocumentFlushed);
+        assert.callOrder(global.promiseDocumentFlushed, pageAction.label.getClientRects);
       });
       it("should set the CSS variable --cfr-label-width correctly", async () => {
         await pageAction.show(fakeRecommendation);