org.unitils.selenium.ScreenshotTakingWebDriver Maven / Gradle / Ivy

package org.unitils.selenium;

import java.io.File;
import java.io.IOException;
import java.lang.reflect.Method;

import org.apache.commons.io.FileUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.junit.Test;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriver;
import org.unitils.testlink.annotation.TestLink;

import com.google.common.base.Supplier;
import com.thoughtworks.selenium.webdriven.WebDriverBackedSelenium;

/**
 * Decorator for the WebDriverBackedSelenium class to handle screenshots.
 *
 * @author Jef Verelst
 *
 * @since 1.0.4
 */
public class ScreenshotTakingWebDriver extends WebDriverBackedSelenium {

    private static final Log LOGGER = LogFactory.getLog(ScreenshotTakingWebDriver.class);

    private int sequenceNr;

    /**
     * Constructor.
     *
     * @param maker : should be of type {@link Supplier}.
     * @param baseUrl : this is the url where the webdriver should start with.
     */
    public ScreenshotTakingWebDriver(Supplier maker, String baseUrl) {
        super(maker, baseUrl);
    }

    /**
     * Constructor.
     *
     * @param baseDriver : should be of type {@link WebDriver}.
     * @param baseUrl : this is the url where the webdriver should start with.
     */
    public ScreenshotTakingWebDriver(WebDriver baseDriver, String baseUrl) {
        super(baseDriver, baseUrl);
    }

    /**
     * Save a screenshot with a default name in the default folder.
     */
    public void saveScreenshot() {
        StringBuilder seqNr = new StringBuilder("" + sequenceNr);
        while (seqNr.length() < 3) {
            seqNr.insert(0, "0");
        }
        this.saveScreenshot("screenshot_" + seqNr.toString());
        sequenceNr++;
    }

    /**
     * Save a screenshot with a specified name in the default folder.
     *
     * @param fileName of the screenshot
     */
    public void saveScreenshot(String fileName) {
        File screenshotDir = new File("./target/screenshot");
        if (!screenshotDir.exists()) {
            screenshotDir.mkdir();
        }
        // apparently the fastes way (except using the sun specific classes, but they are not
        // guaranteed to remain in the JDK, as is the case with a removed reflection class in JDK8).
        StackTraceElement[] st = new Throwable().getStackTrace();
        for (int i = 1; i < st.length; i++) {
            String className = st[i].getClassName();
            String methodName = st[i].getMethodName();
            try {
                Class clazz = Class.forName(className);
                Method method = clazz.getMethod(methodName, new Class[0]);
                if (method.getAnnotation(Test.class) != null) {
                    // we have found the test method... check if there is a @TestLink annotation (if so, use the id), otherwise the name
                    if (method.getAnnotation(TestLink.class) != null) {
                        screenshotDir = new File(screenshotDir, method.getAnnotation(TestLink.class).value());
                    } else {
                        screenshotDir = new File(screenshotDir, methodName);
                    }
                    break;
                }
            } catch (Exception e) {
                // put it in the default screenshot folder
                LOGGER.error(e.getMessage(), e);
            }
        }
        this.saveScreenshot(screenshotDir, fileName);
    }

    /**
     * Store a screenshot with a specified name in a specified folder.
     *
     * @param screenshotDir : the directory where the screenshots should be saved.
     * @param fileName : the name of the new screenshot.
     */
    public void saveScreenshot(File screenshotDir, String fileName) {

        File scrFile = ((TakesScreenshot) this.getWrappedDriver()).getScreenshotAs(OutputType.FILE);
        try {
            FileUtils.copyFile(scrFile, new File(screenshotDir, fileName + ".png"));
            scrFile.deleteOnExit();
        } catch (IOException e) {
            LOGGER.error(e.getMessage(), e);
        }
    }

    /**
     * Create a screenshot when a specified {@link org.openqa.selenium.WebElement} is clicked.
     *
     * @param locator (f.e. a button).
     */
    public void clickWithScreenshot(String locator) {
        // take a screenshot just before we click a button
        this.saveScreenshot();
        super.click(locator);
    }

    /**
     * Create a screenshot when a specified {@link org.openqa.selenium.WebElement} is clicked.
     *
     * @param locator (f.e. a button).
     * @param coordString : the coördinations of that webelement.
     */
    public void clickWithScreenshotAt(String locator, String coordString) {
        // take a screenshot just before we click a button
        this.saveScreenshot();
        super.clickAt(locator, coordString);
    }

    @Override
    public void waitForPageToLoad(String timeout) {
        super.waitForPageToLoad(timeout);
        // take a screenshot just after the page has loaded
        this.saveScreenshot();
    }

    @Override
    public void waitForFrameToLoad(String frameAddress, String timeout) {
        super.waitForFrameToLoad(frameAddress, timeout);
        // take a screenshot just after the page has loaded
        this.saveScreenshot();
    }


}