Java 类org.openqa.selenium.remote.ScreenshotException 实例源码

@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();
    }
  }
}

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();
}
   }

protected String extractScreenShot(final WebDriverException e) {
    final Throwable cause = e.getCause();
    if (cause instanceof ScreenshotException) {
        return ((ScreenshotException) cause).getBase64EncodedScreenshot();
    }
    return null;
}

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;
}

public String extractScreenShot(WebDriverException e) {
    Throwable cause = e.getCause();
    if (cause instanceof ScreenshotException) {
        return ((ScreenshotException) cause).getBase64EncodedScreenshot();
    }
    return null;
}

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;
}

public String extractScreenShot(WebDriverException e) {
  Throwable cause = e.getCause();
  if (cause instanceof ScreenshotException) {
    return ((ScreenshotException) cause).getBase64EncodedScreenshot();
  }
  return null;
}

/**
 * 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;
}

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);
}

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);
    }        
}

@Test
public void testExtractScreenShot2() {
    final ScreenshotException se = Mockito.mock(ScreenshotException.class);
    Mockito.when(se.getBase64EncodedScreenshot()).thenReturn("ex");
    Assert.assertEquals("ex", extractScreenShot(new WebDriverException(se)));
}