--- a/testing/marionette/capture.js
+++ b/testing/marionette/capture.js
@@ -25,82 +25,78 @@ this.capture = {};
  * @param {Array.<Node>=} highlights
  *     Optional array of nodes, around which a border will be marked to
  *     highlight them in the screenshot.
  *
  * @return {HTMLCanvasElement}
  *     The canvas element where the element has been painted on.
  */
 capture.element = function (node, highlights=[]) {
-  let doc = node.ownerDocument;
-  let win = doc.defaultView;
+  let win = node.ownerDocument.defaultView;
   let rect = node.getBoundingClientRect();
 
   return capture.canvas(
-      doc,
+      win,
       rect.left,
       rect.top,
       rect.width,
       rect.height,
       highlights);
 };
 
 /**
- * Take a screenshot of the document's viewport, taking into account
- * the current window's offset.
+ * Take a screenshot of the window's viewport by taking into account
+ * the current offsets.
  *
- * @param {Document} document
- *     The DOM document providing the document element to capture,
- *     and a window for determining the offset of the viewport.
+ * @param {DOMWindow} win
+ *     The DOM window providing the document element to capture,
+ *     and the offsets for the viewport.
  * @param {Array.<Node>=} highlights
  *     Optional array of nodes, around which a border will be marked to
  *     highlight them in the screenshot.
  *
  * @return {HTMLCanvasElement}
  *     The canvas element where the viewport has been painted on.
  */
-capture.viewport = function (document, highlights=[]) {
-  let win = document.defaultView;
-  let docEl = document.documentElement;
+capture.viewport = function (win, highlights=[]) {
+  let rootNode = win.document.documentElement;
 
   return capture.canvas(
-      document,
+      win,
       win.pageXOffset,
       win.pageYOffset,
-      docEl.clientWidth,
-      docEl.clientHeight,
+      rootNode.clientWidth,
+      rootNode.clientHeight,
       highlights);
 };
 
 /**
  * Low-level interface to draw a rectangle off the framebuffer.
  *
- * @param {Document} document
- *     A DOM document providing the window used to the framebuffer,
- *     and interfaces for creating an HTMLCanvasElement.
+ * @param {DOMWindow} win
+ *     The DOM window used for the framebuffer, and providing the interfaces
+ *     for creating an HTMLCanvasElement.
  * @param {number} left
  *     The left, X axis offset of the rectangle.
  * @param {number} top
  *     The top, Y axis offset of the rectangle.
  * @param {number} width
  *     The width dimension of the rectangle to paint.
  * @param {number} height
  *     The height dimension of the rectangle to paint.
  * @param {Array.<Node>=} highlights
  *     Optional array of nodes, around which a border will be marked to
  *     highlight them in the screenshot.
  *
  * @return {HTMLCanvasElement}
  *     The canvas on which the selection from the window's framebuffer
  *     has been painted on.
  */
-capture.canvas = function (document, left, top, width, height, highlights=[]) {
-  let win = document.defaultView;
-
-  let canvas = document.createElementNS(XHTML_NS, "canvas");
+capture.canvas = function (win, left, top, width, height, highlights=[]) {
+  let canvas = win.document.createElementNS(XHTML_NS, "canvas");
   canvas.width = width;
   canvas.height = height;
 
   let ctx = canvas.getContext(CONTEXT_2D);
   ctx.drawWindow(win, left, top, width, height, BG_COLOUR);
   ctx = capture.highlight_(ctx, highlights, top, left);
 
   return canvas;

--- a/testing/marionette/listener.js
+++ b/testing/marionette/listener.js
@@ -1668,17 +1668,17 @@ function screenshot(id, full=true, highl
   let highlightEls = [];
   for (let h of highlights) {
     let el = seenEls.get(h, curContainer);
     highlightEls.push(el);
   }
 
   // viewport
   if (!id && !full) {
-    canvas = capture.viewport(curContainer.frame.document, highlightEls);
+    canvas = capture.viewport(curContainer.frame, highlightEls);
 
   // element or full document element
   } else {
     let node;
     if (id) {
       node = seenEls.get(id, curContainer);
     } else {
       node = curContainer.frame.document.documentElement;