Bug 1606794 - [marionette] Fixed WebDriver:TakeScreenshot to capture frame's content. r=marionette-reviewers,ato

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

--HG--
extra : moz-landing-system : lando

testing/marionette/driver.js

@@ -377,6 +377,33 @@ GeckoDriver.prototype.sendAsync = function(name, data, commandID) {
   });
 };
 
+/**
+ * Get the browsing context.
+ *
+ * @param {boolean=} topContext
+ *     If set to true use the window's top-level browsing context,
+ *     otherwise the one from the currently selected frame. Defaults to false.
+ *
+ * @return {BrowsingContext}
+ *     The browsing context.
+ */
+GeckoDriver.prototype.getBrowsingContext = async function(topContext = false) {
+  let browsingContext = null;
+
+  switch (this.context) {
+    case Context.Chrome:
+      browsingContext = this.getCurrentWindow().docShell.browsingContext;
+      break;
+
+    case Context.Content:
+      const id = await this.listener.getBrowsingContextId(topContext);
+      browsingContext = BrowsingContext.get(id);
+      break;
+  }
+
+  return browsingContext;
+};
+
 /**
  * Get the session's current top-level browsing context.
  *
@@ -3003,13 +3030,9 @@ GeckoDriver.prototype.takeScreenshot = async function(cmd) {
   // Only consider full screenshot if no element has been specified
   full = webEl ? false : full;
 
-  let browsingContext;
   let rect;
-
   switch (this.context) {
     case Context.Chrome:
-      browsingContext = win.docShell.browsingContext;
-
       if (id) {
         let el = this.curBrowser.seenEls.get(webEl, win);
         rect = el.getBoundingClientRect();
@@ -3028,11 +3051,14 @@ GeckoDriver.prototype.takeScreenshot = async function(cmd) {
       break;
 
     case Context.Content:
-      browsingContext = this.curBrowser.contentBrowser.browsingContext;
       rect = await this.listener.getScreenshotRect({ el: webEl, full, scroll });
       break;
   }
 
+  // If no element has been specified use the top-level browsing context.
+  // Otherwise use the browsing context from the currently selected frame.
+  const browsingContext = await this.getBrowsingContext(!webEl);
+
   let canvas = await capture.canvas(
     win,
     browsingContext,

testing/marionette/listener.js

@@ -535,13 +535,14 @@ function dispatch(fn) {
   };
 }
 
-let getPageSourceFn = dispatch(getPageSource);
 let getActiveElementFn = dispatch(getActiveElement);
+let getBrowsingContextIdFn = dispatch(getBrowsingContextId);
 let getElementAttributeFn = dispatch(getElementAttribute);
 let getElementPropertyFn = dispatch(getElementProperty);
 let getElementTextFn = dispatch(getElementText);
 let getElementTagNameFn = dispatch(getElementTagName);
 let getElementRectFn = dispatch(getElementRect);
+let getPageSourceFn = dispatch(getPageSource);
 let getScreenshotRectFn = dispatch(getScreenshotRect);
 let isElementEnabledFn = dispatch(isElementEnabled);
 let findElementContentFn = dispatch(findElementContent);
@@ -577,6 +578,7 @@ function startListeners() {
   addMessageListener("Marionette:findElementContent", findElementContentFn);
   addMessageListener("Marionette:findElementsContent", findElementsContentFn);
   addMessageListener("Marionette:getActiveElement", getActiveElementFn);
+  addMessageListener("Marionette:getBrowsingContextId", getBrowsingContextIdFn);
   addMessageListener("Marionette:getElementAttribute", getElementAttributeFn);
   addMessageListener("Marionette:getElementProperty", getElementPropertyFn);
   addMessageListener("Marionette:getElementRect", getElementRectFn);
@@ -622,6 +624,10 @@ function deregister() {
     findElementsContentFn
   );
   removeMessageListener("Marionette:getActiveElement", getActiveElementFn);
+  removeMessageListener(
+    "Marionette:getBrowsingContextId",
+    getBrowsingContextIdFn
+  );
   removeMessageListener(
     "Marionette:getElementAttribute",
     getElementAttributeFn
@@ -1286,6 +1292,24 @@ function getActiveElement() {
   return evaluate.toJSON(el, seenEls);
 }
 
+/**
+ * Return the current browsing context id.
+ *
+ * @param {boolean=} topContext
+ *     If set to true use the window's top-level browsing context,
+ *     otherwise the one from the currently selected frame. Defaults to false.
+ *
+ * @return {number}
+ *     Id of the browsing context.
+ */
+function getBrowsingContextId(topContext = false) {
+  if (topContext) {
+    return content.docShell.browsingContext.id;
+  }
+
+  return curContainer.frame.docShell.browsingContext.id;
+}
+
 /**
  * Send click event to element.
  *