Bug 1662460 - [marionette] Automatically convert between DOM nodes and element id references in fromJSON() and toJSON(). r=marionette-reviewers,maja_zf

Differential Revision: https://phabricator.services.mozilla.com/D91918

testing/marionette/actors/MarionetteFrameChild.jsm

@@ -12,14 +12,11 @@ const { XPCOMUtils } = ChromeUtils.import(
 
 XPCOMUtils.defineLazyModuleGetters(this, {
   atom: "chrome://marionette/content/atom.js",
-  ContentDOMReference: "resource://gre/modules/ContentDOMReference.jsm",
   element: "chrome://marionette/content/element.js",
   error: "chrome://marionette/content/error.js",
   evaluate: "chrome://marionette/content/evaluate.js",
   interaction: "chrome://marionette/content/interaction.js",
   Log: "chrome://marionette/content/log.js",
-  pprint: "chrome://marionette/content/format.js",
-  WebElement: "chrome://marionette/content/element.js",
 });
 
 XPCOMUtils.defineLazyGetter(this, "logger", () => Log.get());
@@ -121,6 +118,9 @@ class MarionetteFrameChild extends JSWindowActorChild {
           break;
       }
 
+      // The element reference store lives in the parent process. Calling
+      // toJSON() without a second argument here passes element reference ids
+      // of DOM nodes to the parent frame.
       return { data: evaluate.toJSON(result) };
     } catch (e) {
       // Always wrap errors as WebDriverError
@@ -128,83 +128,27 @@ class MarionetteFrameChild extends JSWindowActorChild {
     }
   }
 
-  /**
-   * Wrapper around ContentDOMReference.get with additional steps specific to
-   * Marionette.
-   *
-   * @param {Element} el
-   *     The DOM element to generate the identifier for.
-   *
-   * @return {object} The ContentDOMReference ElementIdentifier for the DOM
-   *     element augmented with a Marionette WebElement reference.
-   */
-  async getElementId(el) {
-    const id = ContentDOMReference.get(el);
-    const webEl = WebElement.from(el);
-    id.webElRef = webEl.toJSON();
-    // Use known WebElement reference if parent process has seen `id` before
-    // TODO - Bug 1666837 - Avoid interprocess element lookup when possible
-    id.webElRef = await this.sendQuery(
-      "MarionetteFrameChild:ElementIdCacheAdd",
-      id
-    );
-    return id;
-  }
-
-  /**
-   * Wrapper around ContentDOMReference.resolve with additional error handling
-   * specific to Marionette.
-   *
-   * @param {ElementIdentifier} id
-   *     The identifier generated via ContentDOMReference.get for a DOM element.
-   *
-   * @return {Element} The DOM element that the identifier was generated for, or
-   *     null if the element does not still exist.
-   *
-   * @throws {StaleElementReferenceError}
-   *     If the element has gone stale, indicating it is no longer
-   *     attached to the DOM, or its node document is no longer the
-   *     active document.
-   */
-  resolveElement(id) {
-    let webEl;
-    if (id.webElRef) {
-      webEl = WebElement.fromJSON(id.webElRef);
-    }
-    const el = ContentDOMReference.resolve(id);
-    if (element.isStale(el, this.contentWindow)) {
-      throw new error.StaleElementReferenceError(
-        pprint`The element reference of ${el || webEl?.uuid} is stale; ` +
-          "either the element is no longer attached to the DOM, " +
-          "it is not in the current frame context, " +
-          "or the document has been refreshed"
-      );
-    }
-    return el;
-  }
-
   // Implementation of WebDriver commands
 
   /** Clear the text of an element.
    *
    * @param {Object} options
-   * @param {ElementIdentifier} options.webEl
+   * @param {Element} options.elem
    */
   clearElement(options = {}) {
-    const { webEl } = options;
-    const el = this.resolveElement(webEl);
-    interaction.clearElement(el);
+    const { elem } = options;
+
+    interaction.clearElement(elem);
   }
 
   /**
    * Click an element.
    */
   async clickElement(options = {}) {
-    const { capabilities, webEl } = options;
-    const el = this.resolveElement(webEl);
+    const { capabilities, elem } = options;
 
     return interaction.clickElement(
-      el,
+      elem,
       capabilities["moz:accessibilityChecks"],
       capabilities["moz:webdriverClick"]
     );
@@ -216,7 +160,7 @@ class MarionetteFrameChild extends JSWindowActorChild {
    *
    * @param {Object} options
    * @param {Object} options.opts
-   * @param {ElementIdentifier} opts.startNode
+   * @param {Element} opts.startNode
    * @param {string} opts.strategy
    * @param {string} opts.selector
    *
@@ -226,13 +170,8 @@ class MarionetteFrameChild extends JSWindowActorChild {
 
     opts.all = false;
 
-    if (opts.startNode) {
-      opts.startNode = this.resolveElement(opts.startNode);
-    }
-
     const container = { frame: this.contentWindow };
-    const el = await element.find(container, strategy, selector, opts);
-    return this.getElementId(el);
+    return element.find(container, strategy, selector, opts);
   }
 
   /**
@@ -241,7 +180,7 @@ class MarionetteFrameChild extends JSWindowActorChild {
    *
    * @param {Object} options
    * @param {Object} options.opts
-   * @param {ElementIdentifier} opts.startNode
+   * @param {Element} opts.startNode
    * @param {string} opts.strategy
    * @param {string} opts.selector
    *
@@ -251,25 +190,20 @@ class MarionetteFrameChild extends JSWindowActorChild {
 
     opts.all = true;
 
-    if (opts.startNode) {
-      opts.startNode = this.resolveElement(opts.startNode);
-    }
-
     const container = { frame: this.contentWindow };
-    const els = await element.find(container, strategy, selector, opts);
-    return Promise.all(els.map(el => this.getElementId(el)));
+    return element.find(container, strategy, selector, opts);
   }
 
   /**
    * Return the active element in the document.
    */
   async getActiveElement() {
-    let el = this.document.activeElement;
-    if (!el) {
+    let elem = this.document.activeElement;
+    if (!elem) {
       throw new error.NoSuchElementError();
     }
 
-    return this.getElementId(el);
+    return elem;
   }
 
   /**
@@ -283,36 +217,33 @@ class MarionetteFrameChild extends JSWindowActorChild {
    * Get the value of an attribute for the given element.
    */
   async getElementAttribute(options = {}) {
-    const { name, webEl } = options;
-    const el = this.resolveElement(webEl);
+    const { name, elem } = options;
 
-    if (element.isBooleanAttribute(el, name)) {
-      if (el.hasAttribute(name)) {
+    if (element.isBooleanAttribute(elem, name)) {
+      if (elem.hasAttribute(name)) {
         return "true";
       }
       return null;
     }
-    return el.getAttribute(name);
+    return elem.getAttribute(name);
   }
 
   /**
    * Get the value of a property for the given element.
    */
   async getElementProperty(options = {}) {
-    const { name, webEl } = options;
-    const el = this.resolveElement(webEl);
+    const { name, elem } = options;
 
-    return typeof el[name] != "undefined" ? el[name] : null;
+    return typeof elem[name] != "undefined" ? elem[name] : null;
   }
 
   /**
    * Get the position and dimensions of the element.
    */
   async getElementRect(options = {}) {
-    const { webEl } = options;
-    const el = this.resolveElement(webEl);
+    const { elem } = options;
 
-    const rect = el.getBoundingClientRect();
+    const rect = elem.getBoundingClientRect();
     return {
       x: rect.x + this.content.pageXOffset,
       y: rect.y + this.content.pageYOffset,
@@ -325,28 +256,27 @@ class MarionetteFrameChild extends JSWindowActorChild {
    * Get the tagName for the given element.
    */
   async getElementTagName(options = {}) {
-    const { webEl } = options;
-    const el = this.resolveElement(webEl);
-    return el.tagName.toLowerCase();
+    const { elem } = options;
+
+    return elem.tagName.toLowerCase();
   }
 
   /**
    * Get the text content for the given element.
    */
   async getElementText(options = {}) {
-    const { webEl } = options;
-    const el = this.resolveElement(webEl);
-    return atom.getElementText(el, this.contentWindow);
+    const { elem } = options;
+
+    return atom.getElementText(elem, this.contentWindow);
   }
 
   /**
    * Get the value of a css property for the given element.
    */
   async getElementValueOfCssProperty(options = {}) {
-    const { name, webEl } = options;
-    const el = this.resolveElement(webEl);
+    const { name, elem } = options;
 
-    const style = this.contentWindow.getComputedStyle(el);
+    const style = this.contentWindow.getComputedStyle(elem);
     return style.getPropertyValue(name);
   }
 
@@ -361,11 +291,10 @@ class MarionetteFrameChild extends JSWindowActorChild {
    * Determine the element displayedness of the given web element.
    */
   async isElementDisplayed(options = {}) {
-    const { capabilities, webEl } = options;
-    const el = this.resolveElement(webEl);
+    const { capabilities, elem } = options;
 
     return interaction.isElementDisplayed(
-      el,
+      elem,
       capabilities["moz:accessibilityChecks"]
     );
   }
@@ -374,11 +303,10 @@ class MarionetteFrameChild extends JSWindowActorChild {
    * Check if element is enabled.
    */
   async isElementEnabled(options = {}) {
-    const { capabilities, webEl } = options;
-    const el = this.resolveElement(webEl);
+    const { capabilities, elem } = options;
 
     return interaction.isElementEnabled(
-      el,
+      elem,
       capabilities["moz:accessibilityChecks"]
     );
   }
@@ -387,11 +315,10 @@ class MarionetteFrameChild extends JSWindowActorChild {
    * Determine whether the referenced element is selected or not.
    */
   async isElementSelected(options = {}) {
-    const { capabilities, webEl } = options;
-    const el = this.resolveElement(webEl);
+    const { capabilities, elem } = options;
 
     return interaction.isElementSelected(
-      el,
+      elem,
       capabilities["moz:accessibilityChecks"]
     );
   }
@@ -400,8 +327,7 @@ class MarionetteFrameChild extends JSWindowActorChild {
    * Send key presses to element after focusing on it.
    */
   async sendKeysToElement(options = {}) {
-    const { capabilities, text, webEl } = options;
-    const el = this.resolveElement(webEl);
+    const { capabilities, elem, text } = options;
 
     const opts = {
       strictFileInteractability: capabilities.strictFileInteractability,
@@ -409,18 +335,17 @@ class MarionetteFrameChild extends JSWindowActorChild {
       webdriverClick: capabilities["moz:webdriverClick"],
     };
 
-    return interaction.sendKeysToElement(el, text, opts);
+    return interaction.sendKeysToElement(elem, text, opts);
   }
 
   /**
    * Switch to the specified frame.
    *
    * @param {Object=} options
-   * @param {(number|ElementIdentifier)=} options.id
-   *     Identifier of the frame to switch to. If it's a number treat it as
-   *     the index for all the existing frames. If it's an ElementIdentifier switch
-   *     to this specific frame. If not specified or `null` switch to the
-   *     top-level browsing context.
+   * @param {(number|Element)=} options.id
+   *     If it's a number treat it as the index for all the existing frames.
+   *     If it's an Element switch to this specific frame.
+   *     If not specified or `null` switch to the top-level browsing context.
    */
   async switchToFrame(options = {}) {
     const { id } = options;
@@ -437,11 +362,9 @@ class MarionetteFrameChild extends JSWindowActorChild {
         );
       }
       browsingContext = childContexts[id];
-      await this.getElementId(browsingContext.embedderElement);
     } else {
-      const frameElement = this.resolveElement(id);
       const context = childContexts.find(context => {
-        return context.embedderElement === frameElement;
+        return context.embedderElement === id;
       });
       if (!context) {
         throw new error.NoSuchFrameError(

testing/marionette/actors/MarionetteFrameParent.jsm

@@ -68,15 +68,13 @@ class MarionetteFrameParent extends JSWindowActorParent {
   // Proxying methods for WebDriver commands
   // TODO: Maybe using a proxy class instead similar to proxy.js
 
-  clearElement(webEl) {
-    return this.sendQuery("MarionetteFrameParent:clearElement", {
-      webEl,
-    });
+  clearElement(elem) {
+    return this.sendQuery("MarionetteFrameParent:clearElement", { elem });
   }
 
-  clickElement(webEl, capabilities) {
+  clickElement(elem, capabilities) {
     return this.sendQuery("MarionetteFrameParent:clickElement", {
-      webEl,
+      elem,
       capabilities,
     });
   }
@@ -105,44 +103,38 @@ class MarionetteFrameParent extends JSWindowActorParent {
     return this.sendQuery("MarionetteFrameParent:getCurrentUrl");
   }
 
-  async getElementAttribute(webEl, name) {
+  async getElementAttribute(elem, name) {
     return this.sendQuery("MarionetteFrameParent:getElementAttribute", {
+      elem,
       name,
-      webEl,
     });
   }
 
-  async getElementProperty(webEl, name) {
+  async getElementProperty(elem, name) {
     return this.sendQuery("MarionetteFrameParent:getElementProperty", {
+      elem,
       name,
-      webEl,
     });
   }
 
-  async getElementRect(webEl) {
-    return this.sendQuery("MarionetteFrameParent:getElementRect", {
-      webEl,
-    });
+  async getElementRect(elem) {
+    return this.sendQuery("MarionetteFrameParent:getElementRect", { elem });
   }
 
-  async getElementTagName(webEl) {
-    return this.sendQuery("MarionetteFrameParent:getElementTagName", {
-      webEl,
-    });
+  async getElementTagName(elem) {
+    return this.sendQuery("MarionetteFrameParent:getElementTagName", { elem });
   }
 
-  async getElementText(webEl) {
-    return this.sendQuery("MarionetteFrameParent:getElementText", {
-      webEl,
-    });
+  async getElementText(elem) {
+    return this.sendQuery("MarionetteFrameParent:getElementText", { elem });
   }
 
-  async getElementValueOfCssProperty(webEl, name) {
+  async getElementValueOfCssProperty(elem, name) {
     return this.sendQuery(
       "MarionetteFrameParent:getElementValueOfCssProperty",
       {
+        elem,
         name,
-        webEl,
       }
     );
   }
@@ -151,32 +143,32 @@ class MarionetteFrameParent extends JSWindowActorParent {
     return this.sendQuery("MarionetteFrameParent:getPageSource");
   }
 
-  async isElementDisplayed(webEl, capabilities) {
+  async isElementDisplayed(elem, capabilities) {
     return this.sendQuery("MarionetteFrameParent:isElementDisplayed", {
-      webEl,
       capabilities,
+      elem,
     });
   }
 
-  async isElementEnabled(webEl, capabilities) {
+  async isElementEnabled(elem, capabilities) {
     return this.sendQuery("MarionetteFrameParent:isElementEnabled", {
-      webEl,
       capabilities,
+      elem,
     });
   }
 
-  async isElementSelected(webEl, capabilities) {
+  async isElementSelected(elem, capabilities) {
     return this.sendQuery("MarionetteFrameParent:isElementSelected", {
-      webEl,
       capabilities,
+      elem,
     });
   }
 
-  async sendKeysToElement(webEl, text, capabilities) {
+  async sendKeysToElement(elem, text, capabilities) {
     return this.sendQuery("MarionetteFrameParent:sendKeysToElement", {
-      webEl,
-      text,
       capabilities,
+      elem,
+      text,
     });
   }
 

testing/marionette/element.js

@@ -19,6 +19,8 @@ const { XPCOMUtils } = ChromeUtils.import(
 );
 
 XPCOMUtils.defineLazyModuleGetters(this, {
+  ContentDOMReference: "resource://gre/modules/ContentDOMReference.jsm",
+
   assert: "chrome://marionette/content/assert.js",
   atom: "chrome://marionette/content/atom.js",
   error: "chrome://marionette/content/error.js",
@@ -764,6 +766,55 @@ element.findClosest = function(startNode, selector) {
   return null;
 };
 
+/**
+ * Wrapper around ContentDOMReference.get with additional steps specific to
+ * Marionette.
+ *
+ * @param {Element} el
+ *     The DOM element to generate the identifier for.
+ *
+ * @return {object} The ContentDOMReference ElementIdentifier for the DOM
+ *     element augmented with a Marionette WebElement reference.
+ */
+element.getElementId = function(el) {
+  const id = ContentDOMReference.get(el);
+  const webEl = WebElement.from(el);
+  id.webElRef = webEl.toJSON();
+  return id;
+};
+
+/**
+ * Wrapper around ContentDOMReference.resolve with additional error handling
+ * specific to Marionette.
+ *
+ * @param {ElementIdentifier} id
+ *     The identifier generated via ContentDOMReference.get for a DOM element.
+ *
+ * @return {Element} The DOM element that the identifier was generated for, or
+ *     null if the element does not still exist.
+ *
+ * @throws {StaleElementReferenceError}
+ *     If the element has gone stale, indicating it is no longer
+ *     attached to the DOM, or its node document is no longer the
+ *     active document.
+ */
+element.resolveElement = function(id) {
+  let webEl;
+  if (id.webElRef) {
+    webEl = WebElement.fromJSON(id.webElRef);
+  }
+  const el = ContentDOMReference.resolve(id);
+  if (element.isStale(el, this.content)) {
+    throw new error.StaleElementReferenceError(
+      pprint`The element reference of ${el || webEl?.uuid} is stale; ` +
+        "either the element is no longer attached to the DOM, " +
+        "it is not in the current frame context, " +
+        "or the document has been refreshed"
+    );
+  }
+  return el;
+};
+
 /**
  * Determines if <var>obj<var> is an HTML or JS collection.
  *

testing/marionette/evaluate.js

@@ -97,7 +97,6 @@ evaluate.sandbox = function(
   } = {}
 ) {
   let unloadHandler;
-
   let marionetteSandbox = sandbox.create(sb.window);
 
   // timeout handler
@@ -179,15 +178,20 @@ evaluate.sandbox = function(
 
 /**
  * Convert any web elements in arbitrary objects to DOM elements by
- * looking them up in the seen element store, or add new ElementIdentifiers to
- * the seen element reference store.
+ * looking them up in the seen element store. For ElementIdentifiers a new
+ * entry in the seen element reference store gets added when running in the
+ * parent process, otherwise ContentDOMReference is used to retrieve the DOM
+ * node.
  *
  * @param {Object} obj
- *     Arbitrary object containing web elements.
+ *     Arbitrary object containing web elements or ElementIdentifiers.
  * @param {(element.Store|element.ReferenceStore)=} seenEls
- *     Known element store to look up web elements from. If `seenEls` is
- *     undefined or an instance of `element.ReferenceStore`, return WebElement.
- *     If `seenEls` is an instance of `element.Store`, return Element.
+ *     Known element store to look up web elements from. If `seenEls` is an
+ *     instance of `element.ReferenceStore`, return WebElement. If `seenEls`
+ *     is an instance of `element.Store`, return Element. If `seenEls` is
+ *     `undefined` the Element from the ContentDOMReference cache is returned
+ *     when executed in the child process, in the parent process the WebElement
+ *     is passed-through.
  * @param {WindowProxy=} win
  *     Current browsing context, if `seenEls` is provided.
  *
@@ -219,19 +223,28 @@ evaluate.fromJSON = function(obj, seenEls = undefined, win = undefined) {
       } else if (Array.isArray(obj)) {
         return obj.map(e => evaluate.fromJSON(e, seenEls, win));
 
-        // web elements
-      } else if (WebElement.isReference(obj)) {
-        let webEl = WebElement.fromJSON(obj);
-        if (seenEls) {
-          return seenEls.get(webEl, win);
-        }
-        return webEl;
-        // ElementIdentifier
-      } else if (
-        seenEls instanceof element.ReferenceStore &&
-        WebElement.isReference(obj.webElRef)
-      ) {
+        // ElementIdentifier and ReferenceStore (used by JSWindowActor)
+      } else if (WebElement.isReference(obj.webElRef)) {
+        if (seenEls instanceof element.ReferenceStore) {
+          // Parent: Store web element reference in the cache
           return seenEls.add(obj);
+        } else if (!seenEls) {
+          // Child: Resolve ElementIdentifier by using ContentDOMReference
+          return element.resolveElement(obj);
+        }
+        throw new TypeError("seenEls is not an instance of ReferenceStore");
+
+        // WebElement and Store (used by framescript)
+      } else if (WebElement.isReference(obj)) {
+        const webEl = WebElement.fromJSON(obj);
+        if (seenEls instanceof element.Store) {
+          // Child: Get web element from the store
+          return seenEls.get(webEl, win);
+        } else if (!seenEls) {
+          // Parent: No conversion. Just return the web element
+          return webEl;
+        }
+        throw new TypeError("seenEls is not an instance of Store");
       }
 
       // arbitrary objects
@@ -254,9 +267,9 @@ evaluate.fromJSON = function(obj, seenEls = undefined, win = undefined) {
  * - Collections, such as `Array<`, `NodeList`, `HTMLCollection`
  *   et al. are expanded to arrays and then recursed.
  *
- * - Elements that are not known web elements are added to the
- *   `seenEls` element store.  Once known, the elements' associated
- *   web element representation is returned.
+ * - Elements that are not known web elements are added to the `seenEls` element
+ *   store, or the ContentDOMReference registry. Once known, the elements'
+ *   associated web element representation is returned.
  *
  * - WebElements are transformed to the corresponding ElementIdentifier
  *   for use in the content process, if an `element.ReferenceStore` is provided.
@@ -303,14 +316,40 @@ evaluate.toJSON = function(obj, seenEls) {
 
     // WebElement
   } else if (WebElement.isReference(obj)) {
+    // Parent: Convert to ElementIdentifier for use in child actor
     if (seenEls instanceof element.ReferenceStore) {
       return seenEls.get(WebElement.fromJSON(obj));
     }
+
     return obj;
 
+    // ElementIdentifier
+  } else if (WebElement.isReference(obj.webElRef)) {
+    // Parent: Pass-through ElementIdentifiers to the child
+    if (seenEls instanceof element.ReferenceStore) {
+      return obj;
+    }
+
+    // Parent: Otherwise return the web element
+    return WebElement.fromJSON(obj.webElRef);
+
     // Element (HTMLElement, SVGElement, XULElement, et al.)
   } else if (element.isElement(obj)) {
+    // Parent
+    if (seenEls instanceof element.ReferenceStore) {
+      throw new TypeError(`ReferenceStore can't be used with Element`);
+
+      // Child: Add element to the Store, return as WebElement
+    } else if (seenEls instanceof element.Store) {
       return seenEls.add(obj);
+    }
+
+    // If no storage has been specified assume we are in a child process.
+    // Evaluation of code will take place in mutable sandboxes, which are
+    // created to waive xrays by default. As such DOM nodes have to be unwaived
+    // before accessing the ownerGlobal is possible, which is needed by
+    // ContentDOMReference.
+    return element.getElementId(Cu.unwaiveXrays(obj));
 
     // custom JSON representation
   } else if (typeof obj.toJSON == "function") {

testing/marionette/test/unit/test_evaluate.js

@@ -14,6 +14,11 @@ class Element {
     this.tagName = tagName;
     this.localName = tagName;
 
+    // Set default properties
+    this.isConnected = true;
+    this.ownerDocument = {};
+    this.ownerGlobal = { document: this.ownerDocument };
+
     for (let attr in attrs) {
       this[attr] = attrs[attr];
     }
@@ -51,11 +56,14 @@ class XULElement extends Element {
 const domEl = new DOMElement("p");
 const svgEl = new SVGElement("rect");
 const xulEl = new XULElement("browser");
+
 const domWebEl = WebElement.from(domEl);
 const svgWebEl = WebElement.from(svgEl);
 const xulWebEl = WebElement.from(xulEl);
+
 const domElId = { id: 1, browsingContextId: 4, webElRef: domWebEl.toJSON() };
 const svgElId = { id: 2, browsingContextId: 5, webElRef: svgWebEl.toJSON() };
+const xulElId = { id: 3, browsingContextId: 6, webElRef: xulWebEl.toJSON() };
 
 const seenEls = new element.Store();
 const elementIdCache = new element.ReferenceStore();
@@ -77,43 +85,6 @@ add_test(function test_toJSON_types() {
   ok(evaluate.toJSON(domEl, seenEls) instanceof WebElement);
   ok(evaluate.toJSON(svgEl, seenEls) instanceof WebElement);
   ok(evaluate.toJSON(xulEl, seenEls) instanceof WebElement);
-  Assert.throws(
-    () => evaluate.toJSON(domEl, elementIdCache),
-    /TypeError/,
-    "Expected ElementIdentifier"
-  );
-  Assert.throws(
-    () => evaluate.toJSON(svgEl, elementIdCache),
-    /TypeError/,
-    "Expected ElementIdentifier"
-  );
-  Assert.throws(
-    () => evaluate.toJSON(xulEl, elementIdCache),
-    /TypeError/,
-    "Expected ElementIdentifier"
-  );
-
-  // WebElement reference, empty elCache
-  Assert.throws(
-    () => evaluate.toJSON(domWebEl, elementIdCache),
-    /NoSuchElementError/
-  );
-  Assert.throws(
-    () => evaluate.toJSON(svgWebEl, elementIdCache),
-    /NoSuchElementError/
-  );
-  Assert.throws(
-    () => evaluate.toJSON(xulWebEl, elementIdCache),
-    /NoSuchElementError/
-  );
-
-  elementIdCache.add(domElId);
-  elementIdCache.add(svgElId);
-
-  deepEqual(evaluate.toJSON(domWebEl, elementIdCache), domElId);
-  deepEqual(evaluate.toJSON(svgWebEl, elementIdCache), svgElId);
-
-  elementIdCache.clear();
 
   // toJSON
   equal(
@@ -131,6 +102,28 @@ add_test(function test_toJSON_types() {
   run_next_test();
 });
 
+add_test(function test_toJSON_types_ReferenceStore() {
+  // Temporarily add custom elements until xpcshell tests
+  // have access to real DOM nodes (including the Window Proxy)
+  elementIdCache.add(domElId);
+  elementIdCache.add(svgElId);
+  elementIdCache.add(xulElId);
+
+  deepEqual(evaluate.toJSON(domWebEl, elementIdCache), domElId);
+  deepEqual(evaluate.toJSON(svgWebEl, elementIdCache), svgElId);
+  deepEqual(evaluate.toJSON(xulWebEl, elementIdCache), xulElId);
+
+  Assert.throws(
+    () => evaluate.toJSON(domEl, elementIdCache),
+    /TypeError/,
+    "Reference store not usable for elements"
+  );
+
+  elementIdCache.clear();
+
+  run_next_test();
+});
+
 add_test(function test_toJSON_sequences() {
   const input = [
     null,
@@ -153,17 +146,10 @@ add_test(function test_toJSON_sequences() {
   equal("foo", actual[4]);
   deepEqual({ bar: "baz" }, actual[5]);
 
-  Assert.throws(
-    () => evaluate.toJSON(input, elementIdCache),
-    /TypeError/,
-    "Expected ElementIdentifier"
-  );
-
   run_next_test();
 });
 
 add_test(function test_toJSON_sequences_ReferenceStore() {
-  elementIdCache.add(domElId);
   const input = [
     null,
     true,
@@ -176,6 +162,15 @@ add_test(function test_toJSON_sequences_ReferenceStore() {
     },
     { bar: "baz" },
   ];
+
+  Assert.throws(
+    () => evaluate.toJSON(input, elementIdCache),
+    /NoSuchElementError/,
+    "Expected no element"
+  );
+
+  elementIdCache.add(domElId);
+
   const actual = evaluate.toJSON(input, elementIdCache);
 
   equal(null, actual[0]);
@@ -196,6 +191,7 @@ add_test(function test_toJSON_objects() {
     boolean: true,
     array: [],
     webElement: domEl,
+    elementId: domElId,
     toJSON: {
       toJSON() {
         return "foo";
@@ -209,26 +205,20 @@ add_test(function test_toJSON_objects() {
   equal(true, actual.boolean);
   deepEqual([], actual.array);
   ok(actual.webElement instanceof WebElement);
+  ok(actual.elementId instanceof WebElement);
   equal("foo", actual.toJSON);
   deepEqual({ bar: "baz" }, actual.object);
 
-  Assert.throws(
-    () => evaluate.toJSON(input, elementIdCache),
-    /TypeError/,
-    "Expected ElementIdentifier"
-  );
-
   run_next_test();
 });
 
 add_test(function test_toJSON_objects_ReferenceStore() {
-  elementIdCache.add(domElId);
-
   const input = {
     null: null,
     boolean: true,
     array: [],
     webElement: domWebEl,
+    elementId: domElId,
     toJSON: {
       toJSON() {
         return "foo";
@@ -236,12 +226,22 @@ add_test(function test_toJSON_objects_ReferenceStore() {
     },
     object: { bar: "baz" },
   };
+
+  Assert.throws(
+    () => evaluate.toJSON(input, elementIdCache),
+    /NoSuchElementError/,
+    "Expected no element"
+  );
+
+  elementIdCache.add(domElId);
+
   const actual = evaluate.toJSON(input, elementIdCache);
 
   equal(null, actual.null);
   equal(true, actual.boolean);
   deepEqual([], actual.array);
   deepEqual(actual.webElement, domElId);
+  deepEqual(actual.elementId, domElId);
   equal("foo", actual.toJSON);
   deepEqual({ bar: "baz" }, actual.object);
 
@@ -251,25 +251,56 @@ add_test(function test_toJSON_objects_ReferenceStore() {
 });
 
 add_test(function test_fromJSON_ReferenceStore() {
-  const input = {
-    id: domElId,
-  };
-  evaluate.fromJSON(input, elementIdCache);
-  deepEqual(elementIdCache.get(domWebEl), domElId);
+  // Add unknown element to reference store
+  let webEl = evaluate.fromJSON(domElId, elementIdCache);
+  deepEqual(webEl, domWebEl);
+  deepEqual(elementIdCache.get(webEl), domElId);
 
+  // Previously seen element is associated with original web element reference
   const domElId2 = {
     id: 1,
     browsingContextId: 4,
     webElRef: WebElement.from(domEl).toJSON(),
   };
-  evaluate.fromJSON(domElId2, elementIdCache);
-  // previously seen element is associated with original web element reference
-  deepEqual(elementIdCache.get(domWebEl), domElId);
+  webEl = evaluate.fromJSON(domElId2, elementIdCache);
+  deepEqual(webEl, domWebEl);
+  deepEqual(elementIdCache.get(webEl), domElId);
+
+  // Store doesn't contain ElementIdentifiers
+  Assert.throws(
+    () => evaluate.fromJSON(domElId, seenEls),
+    /TypeError/,
+    "Expected element.ReferenceStore"
+  );
+
   elementIdCache.clear();
 
   run_next_test();
 });
 
+add_test(function test_fromJSON_Store() {
+  // Pass-through WebElements without adding it to the element store
+  let webEl = evaluate.fromJSON(domWebEl.toJSON());
+  deepEqual(webEl, domWebEl);
+  ok(!seenEls.has(domWebEl));
+
+  // Find element in the element store
+  webEl = seenEls.add(domEl);
+  const el = evaluate.fromJSON(webEl.toJSON(), seenEls);
+  deepEqual(el, domEl);
+
+  // Reference store doesn't contain web elements
+  Assert.throws(
+    () => evaluate.fromJSON(domWebEl.toJSON(), elementIdCache),
+    /TypeError/,
+    "Expected element.Store"
+  );
+
+  seenEls.clear();
+
+  run_next_test();
+});
+
 add_test(function test_isCyclic_noncyclic() {
   for (let type of [true, 42, "foo", [], {}, null, undefined]) {
     ok(!evaluate.isCyclic(type));