import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
@AfterMethod
public void setScreenshot(ITestResult result) {
  if (!result.isSuccess()) {
    try {
      WebDriver returned = new Augmenter().augment(webDriver);
      if (returned != null) {
        File f = ((TakesScreenshot) returned).getScreenshotAs(OutputType.FILE);
        try {
          FileUtils.copyFile(f,
              new File(SCREENSHOT_FOLDER + result.getName() + SCREENSHOT_FORMAT));
        } catch (IOException e) {
          e.printStackTrace();
        }
      }
    } catch (ScreenshotException se) {
      se.printStackTrace();
    }
  }
}
 

import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
protected void saveScreenshot(String name) {
try {
    // WebDriver returned = new Augmenter().augment(webDriver);
    // if (returned != null) {

    File f = ((TakesScreenshot) webDriver).getScreenshotAs(OutputType.FILE);
    screenshotCounter++;
    try {
	FileUtils.copyFile(f, new File(SCREENSHOT_FOLDER + screenshotCounter + "#" + name + "_" + browser.getName() + SCREENSHOT_FORMAT));
    } catch (IOException e) {
	e.printStackTrace();
    }
    // }
} catch (ScreenshotException se) {
    se.printStackTrace();
}
   }
 

import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
protected String extractScreenShot(final WebDriverException e) {
	final Throwable cause = e.getCause();
	if (cause instanceof ScreenshotException) {
		return ((ScreenshotException) cause).getBase64EncodedScreenshot();
	}
	return null;
}
 

import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
public <T> T extractScreenShot(WebDriverException e, OutputType<T> target) {
	if (e.getCause() instanceof ScreenshotException) {
		String base64Str = ((ScreenshotException) e.getCause()).getBase64EncodedScreenshot();
		return target.convertFromBase64Png(base64Str);
	}
	return null;
}
 

import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
public String extractScreenShot(WebDriverException e) {
	Throwable cause = e.getCause();
	if (cause instanceof ScreenshotException) {
		return ((ScreenshotException) cause).getBase64EncodedScreenshot();
	}
	return null;
}
 

import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
private String retrieveScreenshotPathFromException(Throwable originalException, WebDriver driver) {
	if (!this.takeScreenshotOnFailure) {
		return StringUtils.EMPTY;
	}
	try {
		if (originalException instanceof ScreenshotException) {
			return ((ScreenshotException) originalException).getBase64EncodedScreenshot();
		} else if (driver instanceof TakesScreenshot) {
			return ((TakesScreenshot) driver).getScreenshotAs(OutputType.BASE64);
		}
	} catch (Exception se) {
		this.logger.log(Level.FINE, "Failed to retrieve screenshot after failure", se);
	}
	return StringUtils.EMPTY;
}
 

import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
public String extractScreenShot(WebDriverException e) {
  Throwable cause = e.getCause();
  if (cause instanceof ScreenshotException) {
    return ((ScreenshotException) cause).getBase64EncodedScreenshot();
  }
  return null;
}
 

import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
/**
 * Finds screenshot embedded in throwable, if any.
 * @param t exception to search in.
 * @return content of screenshot (if any is present), null otherwise.
 */
public byte[] findScreenshot(Throwable t) {
    byte[] result = null;
    if (t != null) {
        if (t instanceof ScreenshotException) {
            String encodedScreenshot = ((ScreenshotException)t).getBase64EncodedScreenshot();
            result = Base64.getDecoder().decode(encodedScreenshot);
        } else {
            result = findScreenshot(t.getCause());
        }
    }
    return result;
}
 

import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
private void captureExceptionScreenshot(OnomateAssembly.WaitTimeoutException problem) {
    problem.printStackTrace(System.err);
    /*
     * Scan for the screenshot exception
     */
    Throwable previous, next = problem;
    do {
        if( next instanceof ScreenshotException ){
            recordScreenshot((ScreenshotException)next);
        }
        previous = next;
        next = previous.getCause();
    } while(previous != next && next != null);
}
 

import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
private void recordScreenshot(ScreenshotException information){
    String fileName = "screenshot-"+ System.currentTimeMillis()+".png";
    System.out.println("Capturing screenshot to "+fileName);
    String base64 = information.getBase64EncodedScreenshot();
    byte pngData[] = OutputType.BYTES.convertFromBase64Png(base64);
    try {
        FileOutputStream output = new FileOutputStream(fileName, false);
        output.write(pngData);
        output.close();
    } catch (IOException ioe) {
        throw new IllegalStateException("Unable to write screenshot", ioe);
    }        
}
 

import org.openqa.selenium.remote.ScreenshotException; //导入依赖的package包/类
@Test
public void testExtractScreenShot2() {
	final ScreenshotException se = Mockito.mock(ScreenshotException.class);
	Mockito.when(se.getBase64EncodedScreenshot()).thenReturn("ex");
	Assert.assertEquals("ex", extractScreenShot(new WebDriverException(se)));
}