package net.thucydides.core.screenshots;
import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import com.jhlabs.image.BoxBlurFilter;
import net.thucydides.core.guice.Injectors;
import net.thucydides.core.util.EnvironmentVariables;
import net.thucydides.core.webdriver.ProvidedDriverConfiguration;
import net.thucydides.core.webdriver.WebDriverFacade;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.awt.image.ColorModel;
import java.awt.image.WritableRaster;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
/**
* The photographer takes and stores screenshots during the test.
* The actual screenshots are taken using the specified web driver,
* and are stored in the specified target directory. Screenshots
* are numbered sequentially.
*
* @author johnsmart
*/
public class Photographer {
private final WebDriver driver;
private final File targetDirectory;
private Optional<BlurLevel> blurLevel;
private final Logger logger = LoggerFactory.getLogger(Photographer.class);
private ScreenshotProcessor screenshotProcessor;
private EnvironmentVariables environmentVariables;
protected Logger getLogger() {
return logger;
}
public Photographer(final WebDriver driver, final File targetDirectory) {
this(driver, targetDirectory, Injectors.getInjector().getInstance(ScreenshotProcessor.class), null);
}
public Photographer(final WebDriver driver, final File targetDirectory, final BlurLevel blurLevel) {
this(driver, targetDirectory, Injectors.getInjector().getInstance(ScreenshotProcessor.class), blurLevel);
}
public Photographer(final WebDriver driver, final File targetDirectory, final ScreenshotProcessor screenshotProcessor) {
this(driver, targetDirectory, screenshotProcessor, null);
}
public Photographer(final WebDriver driver,
final File targetDirectory,
final ScreenshotProcessor screenshotProcessor,
BlurLevel blurLevel) {
this(driver, targetDirectory, screenshotProcessor, blurLevel,
Injectors.getInjector().getProvider(EnvironmentVariables.class).get() );
}
public Photographer(final WebDriver driver,
final File targetDirectory,
final ScreenshotProcessor screenshotProcessor,
BlurLevel blurLevel,
EnvironmentVariables environmentVariables) {
Preconditions.checkNotNull(targetDirectory);
Preconditions.checkNotNull(screenshotProcessor);
this.driver = driver;
this.targetDirectory = targetDirectory;
this.screenshotProcessor = screenshotProcessor;
this.blurLevel = Optional.fromNullable(blurLevel);
this.environmentVariables = environmentVariables;
}
public Optional<BlurLevel> getBlurLevel() {
return blurLevel;
}
/**
* Take a screenshot of the current browser and store it in the output directory.
*/
public Optional<File> takeScreenshot() {
if (driver != null && driverCanTakeSnapshots()) {
try {
File screenshotTempFile = null;
Object capturedScreenshot = ((TakesScreenshot) driver).getScreenshotAs(OutputType.FILE);
if (isAFile(capturedScreenshot)) {
screenshotTempFile = (File) capturedScreenshot;
} else if (isByteArray(capturedScreenshot)) {
screenshotTempFile = saveScreenshotData((byte[]) capturedScreenshot);
}
if (screenshotTempFile != null && blurLevel.isPresent()) {
screenshotTempFile = blur(screenshotTempFile);
}
if (screenshotTempFile != null) {
String storedFilename = getDigestScreenshotNameFor(screenshotTempFile);
File savedScreenshot = targetScreenshot(storedFilename);
screenshotProcessor.queueScreenshot(new QueuedScreenshot(screenshotTempFile, savedScreenshot));
return Optional.of(savedScreenshot);
}
} catch (Throwable e) {
getLogger().warn("Failed to write screenshot (possibly an out of memory error): " + e.getMessage());
}
}
return Optional.absent();
}
public String getPageSource() {
return driver.getPageSource();
}
private String getDigestScreenshotNameFor(File screenshotTempFile) throws IOException {
ScreenshotDigest screenshotDigest = new ScreenshotDigest(environmentVariables, blurLevel.orNull());
return screenshotDigest.forScreenshot(screenshotTempFile);
}
protected File blur(File srcFile) throws IOException {
BufferedImage srcImage = ImageIO.read(srcFile);
BufferedImage destImage = deepCopy(srcImage);
BoxBlurFilter boxBlurFilter = new BoxBlurFilter();
boxBlurFilter.setRadius(blurLevel.get().getRadius());
boxBlurFilter.setIterations(3);
destImage = boxBlurFilter.filter(srcImage, destImage);
ByteArrayOutputStream outStream = new ByteArrayOutputStream();
ImageIO.write(destImage, "png", outStream);
return saveScreenshotData(outStream.toByteArray());
}
private BufferedImage deepCopy(BufferedImage srcImage) {
ColorModel cm = srcImage.getColorModel();
boolean isAlphaPremultiplied = cm.isAlphaPremultiplied();
WritableRaster raster = srcImage.copyData(null);
return new BufferedImage(cm, raster, isAlphaPremultiplied, null);
}
private File saveScreenshotData(byte[] capturedScreenshot) throws IOException {
Path temporaryScreenshotFile = Files.createTempFile("screenshot", "");
Files.write(temporaryScreenshotFile, capturedScreenshot);
return temporaryScreenshotFile.toFile();
}
private boolean isAFile(Object screenshot) {
return (screenshot instanceof File);
}
private boolean isByteArray(Object screenshot) {
return (screenshot instanceof byte[]);
}
private File targetScreenshot(String storedFilename) {
targetDirectory.mkdirs();
return new File(targetDirectory, storedFilename);
}
protected boolean driverCanTakeSnapshots() {
if (driver == null) {
return false;
} else if (driverIsProvided()) {
ProvidedDriverConfiguration sourceConfig = new ProvidedDriverConfiguration(environmentVariables);
return sourceConfig.getDriverSource().takesScreenshots();
} else if (driver instanceof WebDriverFacade) {
return ((WebDriverFacade) driver).canTakeScreenshots()
&& (((WebDriverFacade) driver).getProxiedDriver() != null);
} else {
return TakesScreenshot.class.isAssignableFrom(driver.getClass());
}
}
private boolean driverIsProvided() {
return false;
}
public void setScreenshotProcessor(ScreenshotProcessor screenshotProcessor) {
this.screenshotProcessor = screenshotProcessor;
}
protected ScreenshotProcessor getScreenshotProcessor() {
return screenshotProcessor;
}
}