CustomizeViewportPastingDecorator.groovy

class ZeroViewportPastingDecorator extends ViewportPastingDecorator {

	private def fixedHeaderHeight
	private def fixedFooterHeight

	ZeroViewportPastingDecorator(ShootingStrategy strategy, int fixedHeaderHeight, int fixedFooterHeight) {
		super(strategy)
		this.fixedHeaderHeight = fixedHeaderHeight
		this.fixedFooterHeight = fixedFooterHeight
	}

	private Coords getShootingCoords(Set<Coords> coords, int pageWidth, int pageHeight, int viewPortHeight) {
		if (coords == null || coords.isEmpty()) {
			return new Coords(0, 0, pageWidth, pageHeight);
		} else {
			return extendShootingArea(Coords.unity(coords), viewPortHeight, pageHeight);
		}
	}

	private Coords extendShootingArea(Coords shootingCoords, int viewportHeight, int pageHeight) {
		int halfViewport = viewportHeight / 2;
		shootingCoords.y = Math.max(shootingCoords.y - halfViewport / 2, 0);
		shootingCoords.height = Math.min(shootingCoords.height + halfViewport, pageHeight);
		return shootingCoords;
	}

	protected void scrollVerticallyInner(JavascriptExecutor js, int scrollY) {
		js.executeScript("scrollTo(0, arguments[0]); return [];", scrollY);
	}

	private void waitForScrolling() {
		try {
			Thread.sleep(scrollTimeout);
		} catch (InterruptedException e) {
			throw new IllegalStateException("Exception while waiting for scrolling", e);
		}
	}

	@Override
	public BufferedImage getScreenshot(WebDriver wd, Set<Coords> coordsSet) {
		JavascriptExecutor js = (JavascriptExecutor) wd;
		int pageHeight = getFullHeight(wd);
		int pageWidth = getFullWidth(wd);
		int viewportHeight = getWindowHeight(wd);
		def shootingArea = getShootingCoords(coordsSet, pageWidth, pageHeight, viewportHeight);

		BufferedImage finalImage = new BufferedImage(pageWidth, (int) shootingArea.height, BufferedImage.TYPE_3BYTE_BGR);
		Graphics2D graphics = finalImage.createGraphics();

		int scrollSize = viewportHeight - fixedHeaderHeight - fixedFooterHeight;
		for (int n = 0;; n++) {
			scrollVerticallyInner(js, (int) shootingArea.y + scrollSize * n);
			waitForScrolling();
			BufferedImage part = getShootingStrategy().getScreenshot(wd);

			if (n == 0) {
				graphics.drawImage(part, 0, (int)(getCurrentScrollY(js) - shootingArea.y), null);
			} else {
				graphics.drawImage(part,
					0,
					(int)(getCurrentScrollY(js) + fixedHeaderHeight),
					pageWidth,
					(int)(getCurrentScrollY(js) + viewportHeight),
					0,
					(int)(fixedHeaderHeight),
					pageWidth,
					(int)(viewportHeight),
					null);
			}

			if (getCurrentScrollY(js) + viewportHeight >= shootingArea.getHeight()) {
				break;
			}
		}

		graphics.dispose();
		return finalImage;
	}

}