Inject client into newly loaded VitalSource book chapters sooner

When a new content frame was found in VitalSource the client was only injected
either if the frame is already loaded, or when the frame next emits a `load`
event. In the latter case this waits until the document and its subresources
have fully loaded. This can be slow in EPUB chapters that have a lot of images.
Improve this by replacing the frame `load` event observer with a call to a new
`onDocumentReady` function which fires as soon as a document becomes interactive
(according to its `readyState`).

 - Rename existing `onDocumentReady` utility to `onNextDocumentReady` to
   make it clear that it only fires once, and change the implementation to
   be a wrapper around a new `onDocumentReady` function.

 - Add new `onDocumentReady` utility which monitors a frame for changes in the
   content document and invokes a callback each time a document becomes _ready_
   (`readyState` is `interactive` or `complete`).

 - Redo the tests for utilities in `frame-observer.js` so that they use
   real iframes rather than fake ones. Iframes have a complex interface
   and loading sequence, so we really need to use the real thing for
   tests to give us confidence.

 - Use `onDocumentReady` in the VitalSource integration to respond more quickly
   to book content loading in a new content frame.

 - Modify several tests for `FrameObserver`, `VitalSourceInjector` and
   `HypothesisInjector` to be less sensitive to exact timings of events,
   as these changed between the previous and new methods for detecting
   when a document is ready.

Fixes https://github.com/hypothesis/client/issues/4270

src/annotator/frame-observer.js

@@ -53,7 +53,7 @@ export class FrameObserver {
   async _addFrame(frame) {
     this._annotatableFrames.add(frame);
     try {
-      await onDocumentReady(frame);
+      await onNextDocumentReady(frame);
       if (this._isDisconnected) {
         return;
       }
@@ -101,44 +101,140 @@ export class FrameObserver {
 }
 
 /**
- * Resolves a Promise when the iframe's document is ready (loaded and parsed)
+ * Test if this is the empty document that a new iframe has before the URL
+ * specified by its `src` attribute loads.
  *
  * @param {HTMLIFrameElement} frame
- * @return {Promise<void>}
- * @throws {Error} if trying to access a document from a cross-origin iframe
  */
-export function onDocumentReady(frame) {
-  return new Promise(resolve => {
-    // @ts-expect-error
-    const frameDocument = frame.contentWindow.document;
-    const { readyState, location } = frameDocument;
-
-    // Web browsers initially load a blank document before the final document.
-    // This blank document is (1) accessible, (2) has an empty body and head,
-    // and (3) has a 'complete' readyState, on Chrome and Safari, and an
-    // 'uninitialized' readyState on Firefox. If a blank document is detected and
-    // there is a 'src' attribute, it is expected that the blank document will be
-    // replaced by the final document.
-    if (
-      location.href === 'about:blank' &&
-      frame.hasAttribute('src') &&
-      frame.src !== 'about:blank'
-    ) {
-      // Listening to `DOMContentLoaded` on the frame's document doesn't work because the
-      // document is replaced. On the other hand, listening the frame's `load`
-      // works because the frame element (as in HTMLIFrameElement) doesn't change.
-      frame.addEventListener('load', () => {
-        resolve();
-      });
+function hasBlankDocumentThatWillNavigate(frame) {
+  return (
+    frame.contentDocument?.location.href === 'about:blank' &&
+    // Do we expect the frame to navigate away from about:blank?
+    frame.hasAttribute('src') &&
+    frame.src !== 'about:blank'
+  );
+}
+
+/**
+ * Wrapper around {@link onDocumentReady} which returns a promise that resolves
+ * the first time that a document in `frame` becomes ready.
+ *
+ * See {@link onDocumentReady} for the definition of _ready_.
+ *
+ * @param {HTMLIFrameElement} frame
+ * @return {Promise<Document>}
+ */
+export function onNextDocumentReady(frame) {
+  return new Promise((resolve, reject) => {
+    const unsubscribe = onDocumentReady(frame, (err, doc) => {
+      unsubscribe();
+      if (doc) {
+        resolve(doc);
+      } else {
+        reject(err);
+      }
+    });
+  });
+}
+
+/**
+ * Register a callback that is invoked when the content document
+ * (`frame.contentDocument`) in a same-origin iframe becomes _ready_.
+ *
+ * A document is _ready_ when its `readyState` is either "interactive" or
+ * "complete". It must also not be the empty document with URL "about:blank"
+ * that iframes have before they navigate to the URL specified by their "src"
+ * attribute.
+ *
+ * The callback is fired both for the document that is in the frame when
+ * `onDocumentReady` is called, as well as for new documents that are
+ * subsequently loaded into the same frame.
+ *
+ * If at any time the frame navigates to an iframe that is cross-origin,
+ * the callback will fire with an error. It will fire again for subsequent
+ * navigations, but due to platform limitations, it will only fire after the
+ * next document fully loads (ie. when the frame's `load` event fires).
+ *
+ * @param {HTMLIFrameElement} frame
+ * @param {(...args: [Error]|[null, Document]) => void} callback
+ * @param {object} options
+ *   @param {number} [options.pollInterval]
+ * @return {() => void} Callback that unsubscribes from future changes
+ */
+export function onDocumentReady(frame, callback, { pollInterval = 10 } = {}) {
+  let pollTimer;
+  let pollForDocumentChange;
+
+  // Visited documents for which we have fired the callback or are waiting
+  // to become ready.
+  const documents = new WeakSet();
+
+  const cancelPoll = () => {
+    clearTimeout(pollTimer);
+    pollTimer = null;
+  };
+
+  // Begin polling for a document change when the current document is about
+  // to go away.
+  const pollOnUnload = () => {
+    if (frame.contentDocument) {
+      frame.contentWindow?.addEventListener('unload', pollForDocumentChange);
+    }
+  };
+
+  const checkForDocumentChange = () => {
+    const currentDocument = frame.contentDocument;
+    if (!currentDocument) {
+      callback(new Error('Frame is cross-origin'));
       return;
     }
 
-    if (readyState === 'loading') {
-      frameDocument.addEventListener('DOMContentLoaded', () => resolve());
+    if (documents.has(currentDocument)) {
       return;
     }
+    documents.add(currentDocument);
+    cancelPoll();
 
-    // State is 'interactive' or 'complete';
-    resolve();
-  });
+    if (!hasBlankDocumentThatWillNavigate(frame)) {
+      const isReady =
+        currentDocument.readyState === 'interactive' ||
+        currentDocument.readyState === 'complete';
+
+      if (isReady) {
+        callback(null, currentDocument);
+      } else {
+        currentDocument.addEventListener('DOMContentLoaded', () =>
+          callback(null, currentDocument)
+        );
+      }
+    }
+
+    // Poll for the next document change.
+    pollOnUnload();
+  };
+
+  pollForDocumentChange = () => {
+    cancelPoll();
+    pollTimer = setInterval(checkForDocumentChange, pollInterval);
+  };
+
+  // Set up observers for signals that the document either has changed or will
+  // soon change. There are two signals with different trade-offs:
+  //
+  //  - Polling after the current document is about to be unloaded. This allows
+  //    us to detect the new document quickly, but may not fire in some
+  //    situations (exact circumstances unclear, but eg. MDN warns about this).
+  //  - The iframe's "load" event. This is guaranteed to fire but only after the
+  //    new document is fully loaded.
+  pollOnUnload();
+  frame.addEventListener('load', checkForDocumentChange);
+
+  // Notify caller about the current document. This fires asynchronously so that
+  // the caller will have received the unsubscribe callback first.
+  setTimeout(() => checkForDocumentChange(), 0);
+
+  return () => {
+    cancelPoll();
+    frame.removeEventListener('load', checkForDocumentChange);
+  };
 }

src/annotator/hypothesis-injector.js

 import { parseJsonConfig } from '../boot/parse-json-config';
 import { generateHexString } from '../shared/random';
-import { onDocumentReady, FrameObserver } from './frame-observer';
+import { onNextDocumentReady, FrameObserver } from './frame-observer';
 
 /**
  * @typedef {import('../types/annotator').Destroyable} Destroyable
@@ -67,7 +67,7 @@ export async function injectClient(frame, config) {
     return;
   }
 
-  await onDocumentReady(frame);
+  await onNextDocumentReady(frame);
 
   // Propagate the client resource locations from the current frame.
   //

src/annotator/integrations/test/vitalsource-test.js

-import { delay } from '../../../test-util/wait';
+import { delay, waitFor } from '../../../test-util/wait';
 import {
   VitalSourceInjector,
   VitalSourceContentIntegration,
@@ -127,7 +127,8 @@ describe('annotator/integrations/vitalsource', () => {
       }, 'Book container element not found');
     });
 
-    it('injects client into content frame', () => {
+    it('injects client into content frame', async () => {
+      await waitFor(() => fakeInjectClient.called);
       assert.calledWith(fakeInjectClient, fakeViewer.contentFrame, fakeConfig);
     });
 
@@ -150,7 +151,7 @@ describe('annotator/integrations/vitalsource', () => {
         assert.notCalled(fakeInjectClient);
 
         fakeViewer.finishChapterLoad(newChapterContent);
-        await delay(0);
+        await waitFor(() => fakeInjectClient.called);
         assert.calledWith(
           fakeInjectClient,
           fakeViewer.contentFrame,
@@ -160,6 +161,7 @@ describe('annotator/integrations/vitalsource', () => {
     });
 
     it("doesn't re-inject if content frame is removed", async () => {
+      await waitFor(() => fakeInjectClient.called);
       fakeInjectClient.resetHistory();
 
       // Remove the content frame. This will trigger a re-injection check, but
@@ -171,6 +173,7 @@ describe('annotator/integrations/vitalsource', () => {
     });
 
     it("doesn't re-inject if content frame siblings change", async () => {
+      await waitFor(() => fakeInjectClient.called);
       fakeInjectClient.resetHistory();
 
       // Modify the DOM tree. This will trigger a re-injection check, but do

src/annotator/integrations/vitalsource.js

 import { ListenerCollection } from '../../shared/listener-collection';
+import { onDocumentReady } from '../frame-observer';
 import { HTMLIntegration } from './html';
 import { preserveScrollPosition } from './html-side-by-side';
 import { ImageTextLayer } from './image-text-layer';
@@ -73,35 +74,6 @@ export class VitalSourceInjector {
     /** @type {WeakSet<HTMLIFrameElement>} */
     const contentFrames = new WeakSet();
 
-    /** @param {HTMLIFrameElement} frame */
-    const injectIfContentReady = frame => {
-      // Check if this frame contains decoded ebook content. If the document has
-      // not yet finished loading, then we rely on this function being called
-      // again once loading completes.
-
-      const body = frame.contentDocument?.body;
-      const isBookContent =
-        body &&
-        // Check that this is not the blank document which is displayed in
-        // brand new iframes before any of their content has loaded.
-        body.children.length > 0 &&
-        // Check that this is not the temporary page containing encrypted and
-        // invisible book content, which is replaced with the real content after
-        // a form submission. These pages look something like:
-        //
-        // ```
-        // <html>
-        //   <title>content</title>
-        //   <body><div id="page-content">{ Base64 encoded data }</div></body>
-        // </html>
-        // ```
-        !body.querySelector('#page-content');
-
-      if (isBookContent) {
-        injectClient(frame, config);
-      }
-    };
-
     const shadowRoot = /** @type {ShadowRoot} */ (bookElement.shadowRoot);
     const injectClientIntoContentFrame = () => {
       const frame = shadowRoot.querySelector('iframe');
@@ -110,10 +82,25 @@ export class VitalSourceInjector {
         return;
       }
       contentFrames.add(frame);
-
-      injectIfContentReady(frame);
-      frame.addEventListener('load', () => {
-        injectIfContentReady(frame);
+      onDocumentReady(frame, (err, document_) => {
+        const body = document_?.body;
+        const isBookContent =
+          body &&
+          // Check that this is not the temporary page containing encrypted and
+          // invisible book content, which is replaced with the real content after
+          // a form submission. These pages look something like:
+          //
+          // ```
+          // <html>
+          //   <title>content</title>
+          //   <body><div id="page-content">{ Base64 encoded data }</div></body>
+          // </html>
+          // ```
+          !body.querySelector('#page-content');
+
+        if (isBookContent) {
+          injectClient(frame, config);
+        }
       });
     };
 

src/annotator/test/frame-observer-test.js

-import { delay } from '../../test-util/wait';
-import { FrameObserver, onDocumentReady, $imports } from '../frame-observer';
+import { delay, waitFor } from '../../test-util/wait';
+import {
+  FrameObserver,
+  onDocumentReady,
+  onNextDocumentReady,
+  $imports,
+} from '../frame-observer';
+
+function waitForEvent(target, event) {
+  return new Promise(resolve => {
+    target.addEventListener(event, () => resolve());
+  });
+}
+
+function waitForCall(spy) {
+  return waitFor(() => spy.called, 300 /* timeout */);
+}
 
 describe('annotator/frame-observer', () => {
+  let container;
+
+  afterEach(() => {
+    container.remove();
+  });
+
+  beforeEach(() => {
+    container = document.createElement('div');
+    document.body.append(container);
+  });
+
   describe('FrameObserver', () => {
-    let container;
     let frameObserver;
     let onFrameAdded;
     let onFrameRemoved;
@@ -26,8 +51,6 @@ describe('annotator/frame-observer', () => {
     }
 
     beforeEach(() => {
-      container = document.createElement('div');
-      document.body.appendChild(container);
       sandbox.stub(console, 'warn');
       $imports.$mock({
         // Disable debouncing
@@ -51,71 +74,48 @@ describe('annotator/frame-observer', () => {
     });
 
     it('triggers onFrameAdded when an annotatable iframe is added', async () => {
-      let iframe = createAnnotatableIFrame();
-      await delay(0);
-
-      assert.calledWith(onFrameAdded, iframe);
-
-      iframe = createAnnotatableIFrame('enable-annotation', 'yes');
-      await delay(0);
-
-      assert.calledWith(onFrameAdded, iframe);
-
-      iframe = createAnnotatableIFrame('enable-annotation', 'true');
-      await delay(0);
-
-      assert.calledWith(onFrameAdded, iframe);
-
-      iframe = createAnnotatableIFrame('enable-annotation', '1');
-      await delay(0);
-
-      assert.calledWith(onFrameAdded, iframe);
-
-      iframe = createAnnotatableIFrame('enable-annotation', 'false'); // the actual value of the attribute is irrelevant
-      await delay(0);
-
+      const iframe = createAnnotatableIFrame();
+      await waitForCall(onFrameAdded);
       assert.calledWith(onFrameAdded, iframe);
     });
 
     it("doesn't trigger onFrameAdded when non-annotatable iframes are added", async () => {
       createAnnotatableIFrame('dummy-attribute');
-      await delay(0);
-
+      await delay(10);
       assert.notCalled(onFrameAdded);
     });
 
     it('removal of the annotatable iframe triggers onFrameRemoved', async () => {
       const iframe = createAnnotatableIFrame();
 
-      await delay(0);
+      await waitForCall(onFrameAdded);
       assert.calledOnce(onFrameAdded);
       assert.calledWith(onFrameAdded, iframe);
 
       iframe.remove();
 
-      await delay(0);
+      await waitForCall(onFrameRemoved);
       assert.calledOnce(onFrameRemoved);
       assert.calledWith(onFrameRemoved, iframe);
     });
 
     it('removal of the `enable-annotation` attribute triggers onFrameRemoved', async () => {
       const iframe = createAnnotatableIFrame();
-      await delay(0);
-
+      await waitForCall(onFrameAdded);
       assert.calledOnce(onFrameAdded);
       assert.calledWith(onFrameAdded, iframe);
 
       iframe.removeAttribute('enable-annotation');
-      await delay(0);
 
+      await waitForCall(onFrameRemoved);
       assert.calledOnce(onFrameRemoved);
       assert.calledWith(onFrameRemoved, iframe);
     });
 
     it('changing the `src` attribute triggers onFrameRemoved', async () => {
       const iframe = createAnnotatableIFrame();
-      await delay(0);
 
+      await waitForCall(onFrameAdded);
       assert.calledOnce(onFrameAdded);
       assert.calledWith(onFrameAdded, iframe);
 
@@ -131,7 +131,7 @@ describe('annotator/frame-observer', () => {
       const iframe = createAnnotatableIFrame();
 
       frameObserver._discoverFrames(); // Emulate a race condition
-      await onDocumentReady(iframe);
+      await onNextDocumentReady(iframe);
 
       assert.notCalled(onFrameAdded);
     });
@@ -140,67 +140,129 @@ describe('annotator/frame-observer', () => {
       const iframe = createAnnotatableIFrame();
       iframe.setAttribute('src', 'http://localhost:1');
 
-      // In this particular case waiting for the FrameObserver to detect the new
-      // iframe may not be enough. Because the browser fetches the URL in `src`
-      // (it is not reachable) it could take longer, that's why, in addition, we
-      // wait for the iframe's document to completely load.
-      await onDocumentReady(iframe);
-      await delay(0);
+      await waitForCall(console.warn);
 
       assert.notCalled(onFrameAdded);
       assert.calledOnce(console.warn);
     });
   });
 
+  function createFrame(src) {
+    const frame = document.createElement('iframe');
+    frame.src = src;
+    container.append(frame);
+    return frame;
+  }
+
+  const sameOriginURL = new URL(
+    '/base/annotator/test/empty.html',
+    document.location.href
+  ).href;
+
+  // A cross-origin local URL that "loads" fast (whether the load succeeds or
+  // fails doesn't matter for these tests). We assume that nothing else is
+  // listening on the port.
+  const crossOriginURL = 'http://localhost:12345/test.html';
+
   describe('onDocumentReady', () => {
-    let fakeIFrame;
-    let fakeIFrameDocument;
-
-    class FakeIFrameDocument extends EventTarget {
-      constructor() {
-        super();
-        this.readyState = 'loading';
-        this.location = {
-          href: 'about:blank',
-        };
-      }
-    }
+    it('invokes callback with current document if it is already ready', async () => {
+      const callback = sinon.stub();
+      const frame = createFrame(sameOriginURL);
+      await waitForEvent(frame, 'load');
+
+      onDocumentReady(frame, callback);
+      await waitForCall(callback);
+
+      assert.calledOnce(callback);
+      assert.calledWith(callback, null);
+      const doc = callback.args[0][1];
+      assert.equal(doc.location.href, sameOriginURL);
+    });
 
-    beforeEach(() => {
-      fakeIFrameDocument = new FakeIFrameDocument();
-      fakeIFrame = document.createElement('div');
-      fakeIFrame.contentWindow = { document: fakeIFrameDocument };
-      fakeIFrame.setAttribute('src', 'http://my.dummy');
+    it('invokes callback when current document becomes ready', async () => {
+      let resolveDocReady;
+      const docReady = new Promise(resolve => (resolveDocReady = resolve));
+      const callback = sinon
+        .stub()
+        .callsFake((err, doc) => resolveDocReady(doc));
+
+      // We use a randomized URL so that the browser won't have a cached copy
+      // that loads instantly. This will force execution through the code path
+      // that waits for 'DOMContentLoaded' before triggering the callback.
+      const docURL = `${sameOriginURL}?random=${Math.random()}`;
+      const frame = createFrame(docURL);
+      onDocumentReady(frame, callback, { pollInterval: 0 });
+      await docReady;
+
+      assert.calledOnce(callback);
+      assert.calledWith(callback, null);
+      const doc = callback.args[0][1];
+      assert.equal(doc.location.href, docURL);
     });
 
-    it('waits for the iframe load event to be triggered if the document is blank', () => {
-      fakeIFrameDocument.location.href = 'about:blank';
-      const onLoad = onDocumentReady(fakeIFrame);
+    it('invokes callback for subsequent navigations to same-origin documents', async () => {
+      const callback = sinon.stub();
+      const frame = createFrame(sameOriginURL);
+      await waitForEvent(frame, 'load');
 
-      // After the initial 'about:blank' document, a new document is loaded.
-      const newDocument = new FakeIFrameDocument();
-      newDocument.location.href = 'http://my.dummy';
-      newDocument.readyState = 'complete';
-      fakeIFrame.contentWindow = { document: newDocument };
-      fakeIFrame.dispatchEvent(new Event('load'));
+      onDocumentReady(frame, callback);
+      await waitForCall(callback);
 
-      return onLoad;
+      frame.src = sameOriginURL + 'v2';
+      await waitForEvent(frame, 'load');
+
+      assert.calledTwice(callback);
     });
 
-    it('waits for the iframe DOMContentLoaded event to be triggered if the document is loading', () => {
-      fakeIFrameDocument.location.href = 'about:srcdoc';
-      fakeIFrameDocument.readyState = 'loading';
-      const onDOMContentLoaded = onDocumentReady(fakeIFrame);
+    it('invokes callback with error if document is cross-origin', async () => {
+      const callback = sinon.stub();
+      const frame = createFrame(crossOriginURL);
+      await waitForEvent(frame, 'load');
 
-      fakeIFrameDocument.dispatchEvent(new Event('DOMContentLoaded'));
+      onDocumentReady(frame, callback);
+      await waitForCall(callback);
 
-      return onDOMContentLoaded;
+      assert.calledOnce(callback);
+      assert.calledWith(callback, sinon.match.instanceOf(Error));
+      const error = callback.args[0][0];
+      assert.equal(error.message, 'Frame is cross-origin');
     });
 
-    it("resolves immediately if document is 'complete' or 'interactive'", () => {
-      fakeIFrameDocument.location.href = 'about:srcdoc';
-      fakeIFrameDocument.readyState = 'complete';
-      return onDocumentReady(fakeIFrame);
+    it('returns a callback that stops polling', async () => {
+      const callback = sinon.stub();
+      const frame = createFrame(sameOriginURL);
+      await waitForEvent(frame, 'load');
+
+      const unsubscribe = onDocumentReady(frame, callback);
+      await waitForCall(callback);
+
+      unsubscribe();
+      frame.src = sameOriginURL + '?v2';
+      await waitForEvent(frame, 'load');
+
+      assert.calledOnce(callback);
+    });
+  });
+
+  describe('onNextDocumentReady', () => {
+    it('it resolves when a document first becomes ready in the frame', async () => {
+      const frame = createFrame(sameOriginURL);
+      const doc = await onNextDocumentReady(frame);
+      assert.equal(doc.location.href, sameOriginURL);
+    });
+
+    it('it rejects if frame document is cross-origin', async () => {
+      const frame = createFrame(crossOriginURL);
+
+      let error;
+      try {
+        await onNextDocumentReady(frame);
+      } catch (e) {
+        error = e;
+      }
+
+      assert.instanceOf(error, Error);
+      assert.equal(error.message, 'Frame is cross-origin');
     });
   });
 });

src/annotator/test/integration/hypothesis-injector-test.js

-import { delay } from '../../../test-util/wait';
-import { DEBOUNCE_WAIT, onDocumentReady } from '../../frame-observer';
+import { delay, waitFor } from '../../../test-util/wait';
+import { DEBOUNCE_WAIT, onNextDocumentReady } from '../../frame-observer';
 import {
   HypothesisInjector,
   injectClient,
@@ -24,8 +24,8 @@ describe('HypothesisInjector integration test', () => {
 
   // Wait for `HypothesisInjector.injectClient` to finish injecting the client
   // into the page.
-  async function waitForInjectClient() {
-    await delay(0);
+  async function waitForInjectClient(frame) {
+    await waitFor(() => getHypothesisScript(frame));
   }
 
   function getHypothesisScript(iframe) {
@@ -125,7 +125,7 @@ describe('HypothesisInjector integration test', () => {
     // Now initialize
     createHypothesisInjector();
 
-    await waitForInjectClient();
+    await waitForInjectClient(validFrame);
 
     assert.isNotNull(
       getHypothesisScript(validFrame),
@@ -142,7 +142,7 @@ describe('HypothesisInjector integration test', () => {
     const iframe = createAnnotatableIFrame();
 
     createHypothesisInjector();
-    await waitForInjectClient();
+    await waitForInjectClient(iframe);
 
     const scriptElement = getHypothesisScript(iframe);
     assert.isNotNull(
@@ -165,7 +165,7 @@ describe('HypothesisInjector integration test', () => {
     iframe.contentDocument.body.append(configScript);
 
     createHypothesisInjector();
-    await onDocumentReady(iframe);
+    await onNextDocumentReady(iframe);
 
     assert.isNull(
       getHypothesisScript(iframe),
@@ -181,7 +181,7 @@ describe('HypothesisInjector integration test', () => {
     const iframe = createAnnotatableIFrame();
 
     await waitForFrameObserver();
-    await waitForInjectClient();
+    await waitForInjectClient(iframe);
     assert.isNotNull(
       getHypothesisScript(iframe),
       'expected dynamically added iframe to include the Hypothesis script'
@@ -193,7 +193,7 @@ describe('HypothesisInjector integration test', () => {
 
     // Now initialize
     createHypothesisInjector();
-    await waitForInjectClient();
+    await waitForInjectClient(iframe);
 
     assert.isNotNull(
       getHypothesisScript(iframe),
@@ -207,7 +207,7 @@ describe('HypothesisInjector integration test', () => {
     assert.isNull(getHypothesisScript(iframe));
 
     await waitForFrameObserver();
-    await waitForInjectClient();
+    await waitForInjectClient(iframe);
 
     assert.isNotNull(
       getHypothesisScript(iframe),
@@ -223,7 +223,7 @@ describe('HypothesisInjector integration test', () => {
     const iframe = createAnnotatableIFrame();
 
     await waitForFrameObserver();
-    await waitForInjectClient();
+    await waitForInjectClient(iframe);
 
     assert.isNotNull(
       getHypothesisScript(iframe),
@@ -237,7 +237,7 @@ describe('HypothesisInjector integration test', () => {
     assert.isNull(getHypothesisScript(iframe));
 
     await waitForFrameObserver();
-    await waitForInjectClient();
+    await waitForInjectClient(iframe);
 
     assert.isNotNull(
       getHypothesisScript(iframe),