SimonStewart: Applying patch from TerenceHaddock that enables screens…

…hots to be saved using a nice API. This also includes the Base64Encoder from XStream, because it's painful to add a library dependency for base64 encoding. *sigh* We may need to reimplement this ourselves at a later date.

r7661

CREDITS.txt

@@ -34,4 +34,4 @@ Michele Sama
 Kenneth Leftin
 Darrell Deboer
 Muthu Kannan
-
+Terence Haddock

common/src/java/org/openqa/selenium/OutputType.java

@@ -0,0 +1,88 @@
+/*
+Copyright 2007-2009 WebDriver committers
+Copyright 2007-2009 Google Inc.
+Portions copyright 2007 ThoughtWorks, Inc
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package org.openqa.selenium;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+
+import org.openqa.selenium.internal.Base64Encoder;
+
+/**
+ * Defines the output type for a screenshot. See org.openqa.selenium.Screenshot
+ * for usage and examples.
+ * 
+ * @see TakesScreenshot
+ * @param <T> Type for the screenshot output.
+ */
+public interface OutputType<T> {
+  /**
+   * Obtain the screenshot as base64 data.
+   */
+  OutputType<String> BASE64 = new OutputType<String>() {
+  	public String convertFromBase64Png(String base64Png) {
+      return base64Png;
+  	}
+  };
+
+  /**
+   * Obtain the screenshot as raw bytes.
+   */
+  OutputType<byte[]> BYTES = new OutputType<byte[]>() {
+  	public byte[] convertFromBase64Png(String base64Png) {
+		  return new Base64Encoder().decode(base64Png);
+  	} 
+  };
+
+  /**
+   * Obtain the screenshot into a temporary file that will be deleted once the
+   * JVM exits. It is up to users to make a copy of this file.
+   */
+  OutputType<File> FILE = new OutputType<File>() {
+  	public File convertFromBase64Png(String base64Png) {
+      FileOutputStream fos = null;
+  		try {
+  			byte[] data = BYTES.convertFromBase64Png(base64Png);
+  			File tmpFile = File.createTempFile("screenshot", ".png");
+        tmpFile.deleteOnExit();
+  			fos = new FileOutputStream(tmpFile);
+  			fos.write(data);
+  			return tmpFile;
+  		} catch (IOException e) {
+  			throw new WebDriverException(e);
+  		} finally {
+        if (fos != null) {
+          try {
+            fos.close();
+          } catch (IOException e) {
+            // Nothing sensible to do
+          }
+        }
+      }
+  	}
+  };
+
+  /**
+   * Convert the given base64 png to a requested format.
+   *
+   * @param base64Png base64 encoded png.
+   * @return png encoded into requested format.
+   */
+  T convertFromBase64Png(String base64Png);
+}

common/src/java/org/openqa/selenium/TakesScreenshot.java

@@ -0,0 +1,44 @@
+/*
+Copyright 2007-2009 WebDriver committers
+Copyright 2007-2009 Google Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package org.openqa.selenium;
+
+
+/**
+ * Indicates a driver that can capture a screenshot and store it in different
+ * ways.
+ * <p>
+ * Example usage:
+ * <pre>
+ * import static openqa.selenium.OutputType.*;
+ * 
+ * File screenshotFile = ((Screenshot)driver).getScreenshotAs(file);
+ * String screenshotBase64 = ((Screenshot)driver).getScreenshotAs(base64);
+ * </pre>
+ * 
+ * @see OutputType
+ */
+public interface TakesScreenshot {
+  /**
+   * Capture the screenshot and store it in the specified location.
+   *
+   * @param <X> Return type for getScreenshotAs.
+   * @param target target type, @see OutputType
+   * @return Object in which is stored information about the screenshot.
+   * @throws WebDriverException on failure.
+   */
+  <X> X getScreenshotAs(OutputType<X> target) throws WebDriverException;
+}

common/src/java/org/openqa/selenium/internal/Base64Encoder.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2004 Joe Walnes.
+ * Copyright (C) 2006, 2007 XStream Committers.
+ * All rights reserved.
+ *
+ * The software in this package is published under the terms of the BSD
+ * style license a copy of which has been included with this distribution in
+ * the LICENSE.txt file.
+ *
+ * Created on 06. August 2004 by Joe Walnes
+ */
+// original package name: com.thoughtworks.xstream.core.util
+package org.openqa.selenium.internal;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.Reader;
+import java.io.StringReader;
+
+/**
+ * Encodes binary data to plain text as Base64.
+ *
+ * <p>Despite there being a gazillion other Base64 implementations out there, this has been written as part of XStream as
+ * it forms a core part but is too trivial to warrant an extra dependency.</p>
+ *
+ * <p>This meets the standard as described in RFC 1521, section 5.2 <http://www.freesoft.org/CIE/RFC/1521/7.htm>, allowing
+ * other Base64 tools to manipulate the data.</p>
+ *
+ * @author Joe Walnes
+ */
+public class Base64Encoder {
+
+    // Here's how encoding works:
+    //
+    // 1) Incoming bytes are broken up into groups of 3 (each byte having 8 bits).
+    //
+    // 2) The combined 24 bits (3 * 8) are split into 4 groups of 6 bits.
+    //
+    // input  |------||------||------| (3 values each with 8 bits)
+    //        101010101010101010101010
+    // output |----||----||----||----| (4 values each with 6 bits)
+    //
+    // 3) Each of these 4 groups of 6 bits are converted back to a number, which will fall in the range of 0 - 63.
+    //
+    // 4) Each of these 4 numbers are converted to an alphanumeric char in a specified mapping table, to create
+    //    a 4 character string.
+    //
+    // 5) This is repeated for all groups of three bytes.
+    //
+    // 6) Special padding is done at the end of the stream using the '=' char.
+
+    private static final char[] SIXTY_FOUR_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/".toCharArray();
+    private static final int[] REVERSE_MAPPING = new int[123];
+
+    static {
+        for (int i = 0; i < SIXTY_FOUR_CHARS.length; i++) REVERSE_MAPPING[SIXTY_FOUR_CHARS[i]] = i + 1;
+    }
+
+    public String encode(byte[] input) {
+        StringBuffer result = new StringBuffer();
+        int outputCharCount = 0;
+        for (int i = 0; i < input.length; i += 3) {
+            int remaining = Math.min(3, input.length - i);
+            int oneBigNumber = (input[i] & 0xff) << 16 | (remaining <= 1 ? 0 : input[i + 1] & 0xff) << 8 | (remaining <= 2 ? 0 : input[i + 2] & 0xff);
+            for (int j = 0; j < 4; j++) result.append(remaining + 1 > j ? SIXTY_FOUR_CHARS[0x3f & oneBigNumber >> 6 * (3 - j)] : '=');
+            if ((outputCharCount += 4) % 76 == 0) result.append('\n');
+        }
+        return result.toString();
+    }
+
+    public byte[] decode(String input) {
+        try {
+            ByteArrayOutputStream out = new ByteArrayOutputStream();
+            StringReader in = new StringReader(input);
+            for (int i = 0; i < input.length(); i += 4) {
+                int a[] = {mapCharToInt(in), mapCharToInt(in), mapCharToInt(in), mapCharToInt(in)};
+                int oneBigNumber = (a[0] & 0x3f) << 18 | (a[1] & 0x3f) << 12 | (a[2] & 0x3f) << 6 | (a[3] & 0x3f);
+                for (int j = 0; j < 3; j++) if (a[j + 1] >= 0) out.write(0xff & oneBigNumber >> 8 * (2 - j));
+            }
+            return out.toByteArray();
+        } catch (IOException e) {
+            throw new Error(e + ": " + e.getMessage());
+        }
+    }
+
+    private int mapCharToInt(Reader input) throws IOException {
+        int c;
+        while ((c = input.read()) != -1) {
+            int result = REVERSE_MAPPING[c];
+            if (result != 0) return result -1;
+            if (c == '=') return -1;
+        }
+        return -1;
+    }
+}

common/test/java/org/openqa/selenium/OutputTypeTest.java

@@ -0,0 +1,45 @@
+/*
+Copyright 2009 WebDriver committers
+Copyright 2009 Google Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package org.openqa.selenium;
+
+import java.io.File;
+
+public class OutputTypeTest extends AbstractDriverTestCase {
+	public static final String TEST_BASE64 = "ABADABAD";
+	public static final byte[] TEST_BYTES = new byte[]{ 0, 16, 3, 0, 16, 3 };
+
+	public void testBase64() {
+		assertEquals(TEST_BASE64, OutputType.BASE64.convertFromBase64Png(TEST_BASE64));
+	}
+
+	public void testBytes() {
+		byte[] bytes = OutputType.BYTES
+        .convertFromBase64Png(TEST_BASE64);
+		assertEquals(TEST_BYTES.length, bytes.length);
+		for (int i = 0; i < TEST_BYTES.length; i++) {
+			assertEquals("index " + i, TEST_BYTES[i], bytes[i]);
+		}
+	}
+
+	public void testFiles() {
+		File tmpFile = OutputType.FILE
+        .convertFromBase64Png(TEST_BASE64);
+		assertTrue(tmpFile.exists());
+		assertEquals(TEST_BYTES.length, tmpFile.length());
+		tmpFile.delete();
+	}
+}

common/test/java/org/openqa/selenium/TakesScreenshotTest.java

@@ -0,0 +1,52 @@
+/*
+Copyright 2009 WebDriver committers
+Copyright 2009 Google Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package org.openqa.selenium;
+
+import static org.openqa.selenium.OutputType.*;
+
+import java.io.File;
+
+public class TakesScreenshotTest extends AbstractDriverTestCase {
+  public void testSaveScreenshotAsFile() throws Exception {
+    if (!isAbleToTakeScreenshots(driver)) {
+      return;
+    }
+
+    File tempFile = getScreenshot().getScreenshotAs(OutputType.FILE);
+    assertTrue(tempFile.exists());
+    assertTrue(tempFile.length() > 0);
+    tempFile.delete();
+  }
+
+  public void testCaptureToBase64() throws Exception {
+    if (!isAbleToTakeScreenshots(driver)) {
+      return;
+    }
+
+    String screenshot = getScreenshot().getScreenshotAs(BASE64);
+    assertTrue(screenshot.length() > 0);
+  }
+
+  public TakesScreenshot getScreenshot() {
+    return (TakesScreenshot)driver;
+  }
+
+  private boolean isAbleToTakeScreenshots(WebDriver driver) throws Exception {
+    return driver instanceof TakesScreenshot;
+  }
+}

firefox/src/extension/components/firefoxDriver.js

@@ -725,3 +725,17 @@ FirefoxDriver.prototype.saveScreenshot = function(respond, pngFile) {
   }
   respond.send();
 };
+
+
+FirefoxDriver.prototype.getScreenshotAsBase64 = function(respond) {
+  var window = Utils.getBrowser(respond.context).contentWindow;
+  try {
+    var canvas = Screenshooter.grab(window);
+    respond.isError = false;
+    respond.response = Screenshooter.toBase64(canvas);
+  } catch (e) {
+    respond.isError = true;
+    respond.response = 'Could not take screenshot of current page - ' + e;
+  }
+  respond.send();
+};

firefox/src/extension/components/screenshooter.js

@@ -42,6 +42,17 @@ Screenshooter.grab = function(window) {
 };
 
 
+Screenshooter.toBase64 = function(canvas) {
+  var dataUrl = canvas.toDataURL('image/png');
+  var index = dataUrl.indexOf('base64,');
+  if (index == -1) {
+    // No base64 data marker.
+    throw new Error("Invalid base64 data: " + dataUrl);
+  }
+  return dataUrl.substring(index + 'base64,'.length);
+};
+
+
 Screenshooter.save = function(canvas, filepath) {
   var cc = Components.classes;
   var ci = Components.interfaces;