无法使用HtmlUnitDriver进行screenShot [Selenium WebDriver java]

Question

I am wanting to take a screenshot of a page using HtmlUnitDriver I came across this Link where this guy has made a custom HTML unit driver to take the screenshot. But unfortunately, while implementing that I am getting an exception.

"Exception in thread "main" java.lang.ClassCastException: [B cannot be cast to java.io.File at Test.main(Test.java:39)"

My code is as follows-

import java.io.File;
import java.io.IOException;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.WebDriver;
import com.gargoylesoftware.htmlunit.BrowserVersion;

public class Test extends ScreenCaptureHtmlUnitDriver {

    public static void main(String[] args) throws InterruptedException, IOException {

        WebDriver driver = new ScreenCaptureHtmlUnitDriver(BrowserVersion.FIREFOX_38);
        driver.get("https://www.google.com/?gws_rd=ssl");
        try{
        File scrFile = ((ScreenCaptureHtmlUnitDriver) driver).getScreenshotAs(OutputType.FILE);
        FileUtils.copyFile(scrFile, new File("D:\\TEMP.PNG"));
        }catch (Exception e) {
            e.printStackTrace();
        }
    }
}

HtmlUnit driver which I am using(the one which is in the link) is this-

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.URL;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;
import org.apache.commons.io.FilenameUtils;
import org.apache.commons.io.IOUtils;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriverException;
import org.openqa.selenium.htmlunit.HtmlUnitDriver;
import org.openqa.selenium.internal.Base64Encoder;
import org.openqa.selenium.remote.CapabilityType;
import org.openqa.selenium.remote.DesiredCapabilities;
import com.gargoylesoftware.htmlunit.BrowserVersion;
import com.gargoylesoftware.htmlunit.WebClient;
import com.gargoylesoftware.htmlunit.WebRequest;
import com.gargoylesoftware.htmlunit.WebWindow;
import com.gargoylesoftware.htmlunit.html.HtmlElement;
import com.gargoylesoftware.htmlunit.html.HtmlPage;

public class ScreenCaptureHtmlUnitDriver extends HtmlUnitDriver implements TakesScreenshot {

private static Map<String, byte[]> imagesCache = Collections.synchronizedMap(new HashMap<String, byte[]>());

private static Map<String, String> cssjsCache = Collections.synchronizedMap(new HashMap<String, String>());

// http://stackoverflow.com/questions/4652777/java-regex-to-get-the-urls-from-css
private final static Pattern cssUrlPattern = Pattern.compile("background(-image)?[\\s]*:[^url]*url[\\s]*\\([\\s]*([^\\)]*)[\\s]*\\)[\\s]*");// ?<url>

public ScreenCaptureHtmlUnitDriver() {
    super();
}

public ScreenCaptureHtmlUnitDriver(boolean enableJavascript) {
    super(enableJavascript);
}

public ScreenCaptureHtmlUnitDriver(Capabilities capabilities) {
    super(capabilities);
}

public ScreenCaptureHtmlUnitDriver(BrowserVersion version) {
    super(version);
    DesiredCapabilities var = ((DesiredCapabilities) getCapabilities());
    var.setCapability(CapabilityType.TAKES_SCREENSHOT, true);
}

//@Override
@SuppressWarnings("unchecked")
public <X> X getScreenshotAs(OutputType<X> target) throws WebDriverException {
    byte[] archive = new byte[0];
    try {
        archive = downloadCssAndImages(getWebClient(), (HtmlPage) getCurrentWindow().getEnclosedPage());
    } catch (Exception e) {
    }
    if(target.equals(OutputType.BASE64)){
        return target.convertFromBase64Png(new Base64Encoder().encode(archive));
    }
    if(target.equals(OutputType.BYTES)){
        return (X) archive;
    }
    return (X) archive;
}

// http://stackoverflow.com/questions/2244272/how-can-i-tell-htmlunits-webclient-to-download-images-and-css
protected byte[] downloadCssAndImages(WebClient webClient, HtmlPage page) throws Exception {
    WebWindow currentWindow = webClient.getCurrentWindow();
    Map<String, String> urlMapping = new HashMap<String, String>();
    Map<String, byte[]> files = new HashMap<String, byte[]>();
    WebWindow window = null;
    try {
        window = webClient.getWebWindowByName(page.getUrl().toString()+"_screenshot");
        webClient.getPage(window, new WebRequest(page.getUrl()));
    } catch (Exception e) {
        window = webClient.openWindow(page.getUrl(), page.getUrl().toString()+"_screenshot");
    }

    String xPathExpression = "//*[name() = 'img' or name() = 'link' and (@type = 'text/css' or @type = 'image/x-icon') or  @type = 'text/javascript']";
    List<?> resultList = page.getByXPath(xPathExpression);

    Iterator<?> i = resultList.iterator();
    while (i.hasNext()) {
        try {
            HtmlElement el = (HtmlElement) i.next();
            String resourceSourcePath = el.getAttribute("src").equals("") ? el.getAttribute("href") : el
                    .getAttribute("src");
            if (resourceSourcePath == null || resourceSourcePath.equals(""))
                continue;
            URL resourceRemoteLink = page.getFullyQualifiedUrl(resourceSourcePath);
            String resourceLocalPath = mapLocalUrl(page, resourceRemoteLink, resourceSourcePath, urlMapping);
            urlMapping.put(resourceSourcePath, resourceLocalPath);
            if (!resourceRemoteLink.toString().endsWith(".css")) {
                byte[] image = downloadImage(webClient, window,  resourceRemoteLink);
                files.put(resourceLocalPath, image);
            } else {
                String css = downloadCss(webClient, window, resourceRemoteLink);
                for (String cssImagePath : getLinksFromCss(css)) {
                    URL cssImagelink = page.getFullyQualifiedUrl(cssImagePath.replace("\"", "").replace("\'", "")
                            .replace(" ", ""));
                    String cssImageLocalPath = mapLocalUrl(page, cssImagelink, cssImagePath, urlMapping);
                    files.put(cssImageLocalPath, downloadImage(webClient, window, cssImagelink));
                }
                files.put(resourceLocalPath, replaceRemoteUrlsWithLocal(css, urlMapping)
                        .replace("resources/", "./").getBytes());
            }
        } catch (Exception e) {
        }
    }
    String pagesrc =  replaceRemoteUrlsWithLocal(page.getWebResponse().getContentAsString(), urlMapping);
    files.put("page.html", pagesrc.getBytes());
    webClient.setCurrentWindow(currentWindow);
    return createZip(files);
}

String downloadCss(WebClient webClient, WebWindow window, URL resourceUrl) throws Exception {
    if (cssjsCache.get(resourceUrl.toString()) == null) {
        cssjsCache.put(resourceUrl.toString(), webClient.getPage(window, new  WebRequest(resourceUrl))
                .getWebResponse().getContentAsString());

    }
    return cssjsCache.get(resourceUrl.toString());
}

byte[] downloadImage(WebClient webClient, WebWindow window, URL resourceUrl)  throws Exception {
    if (imagesCache.get(resourceUrl.toString()) == null) {
        imagesCache.put(
                resourceUrl.toString(),
                IOUtils.toByteArray(webClient.getPage(window, new  WebRequest(resourceUrl)).getWebResponse()
                        .getContentAsStream()));
    }
    return imagesCache.get(resourceUrl.toString());
}

 public static byte[] createZip(Map<String, byte[]> files) throws IOException      {
    ByteArrayOutputStream bos = new ByteArrayOutputStream();
    ZipOutputStream zipfile = new ZipOutputStream(bos);
    Iterator<String> i = files.keySet().iterator();
    String fileName = null;
    ZipEntry zipentry = null;
    while (i.hasNext()) {
        fileName = i.next();
        zipentry = new ZipEntry(fileName);
        zipfile.putNextEntry(zipentry);
        zipfile.write(files.get(fileName));
    }
    zipfile.close();
    return bos.toByteArray();
}

    List<String> getLinksFromCss(String css) {
    List<String> result = new LinkedList<String>();
    Matcher m = cssUrlPattern.matcher(css);
    while (m.find()) { // find next match
        result.add( m.group(2));
    }
    return result;
}

 String replaceRemoteUrlsWithLocal(String source, Map<String, String>  replacement) {
    for (String object : replacement.keySet()) {
        // background:url(http://org.com/images/image.gif)
        source = source.replace(object, replacement.get(object));
    }
    return source;
}

String mapLocalUrl(HtmlPage page, URL link, String path, Map<String, String>  replacementToAdd) throws Exception {
    String resultingFileName = "resources/" +    FilenameUtils.getName(link.getFile());
    replacementToAdd.put(path, resultingFileName);
    return resultingFileName;
}

}

UPDATE

Code provided by Andrew works- but I wanted to know if there is a way by which we can download only selected resources. For eg this website I would like to download only captcha image those id is "//*[@id='cimage']" because downloading all the resources will take a long time. Is there a way by which we can download only the specific resource. Because with the existing code provided below all the resources get downloaded.

byte[] zipFileBytes = ((ScreenCaptureHtmlUnitDriver) driver).getScreenshotAs(OutputType.BYTES);
FileUtils.writeByteArrayToFile(new File("D:\\TEMP.PNG"), zipFileBytes);

Answer 1

The error says that the code is trying to convert a byte[] to a File . It's easy to see why if you just strip out the unused paths from getScreenshotAs :

public <X> X getScreenshotAs(OutputType<X> target) throws WebDriverException {
    byte[] archive = new byte[0];
    try {
        archive = downloadCssAndImages(getWebClient(), (HtmlPage) getCurrentWindow().getEnclosedPage());
    } catch (Exception e) {
    }
    return (X) archive;
}

There's no way you can get a File out of that. OutputType.FILE is not supported, so you have to handle file output yourself. Luckily, that's easy. You can change your code to:

byte[] zipFileBytes = ((ScreenCaptureHtmlUnitDriver) driver).getScreenshotAs(OutputType.BYTES);
FileUtils.writeByteArrayToFile(new File("D:\\TEMP.PNG"), zipFileBytes);

See FileUtils.writeByteArrayToFile () for more.

Answer 2

Check this out this may helpful for you

File scrFile = ((TakesScreenshot)driver).getScreenshotAs(OutputType.FILE);
FileUtils.copyFile(scrFile, new File("C:/Users/home/Desktop/screenshot.png"));// copy it somewhere