From db34e828386f9c1e66012fa32cf3a86e2052c071 Mon Sep 17 00:00:00 2001 From: Uladzislau Arlouski Date: Fri, 20 Dec 2024 17:24:35 +0200 Subject: [PATCH] Add shooting strategy to take screenshot using CDP --- pom.xml | 8 ++- .../pazone/ashot/CdpShootingStrategy.java | 67 +++++++++++++++++++ .../pazone/ashot/CdpShootingStrategyTest.java | 66 ++++++++++++++++++ 3 files changed, 140 insertions(+), 1 deletion(-) create mode 100644 src/main/java/pazone/ashot/CdpShootingStrategy.java create mode 100644 src/test/java/pazone/ashot/CdpShootingStrategyTest.java diff --git a/pom.xml b/pom.xml index 3fabd1b..bd81281 100644 --- a/pom.xml +++ b/pom.xml @@ -27,6 +27,7 @@ UTF-8 + 4.27.0 AShot WebDriver Utility @@ -35,7 +36,12 @@ org.seleniumhq.selenium selenium-remote-driver - 4.27.0 + ${selenium.version} + + + org.seleniumhq.selenium + selenium-chromium-driver + ${selenium.version} commons-io diff --git a/src/main/java/pazone/ashot/CdpShootingStrategy.java b/src/main/java/pazone/ashot/CdpShootingStrategy.java new file mode 100644 index 0000000..b35852c --- /dev/null +++ b/src/main/java/pazone/ashot/CdpShootingStrategy.java @@ -0,0 +1,67 @@ +package pazone.ashot; + +import java.awt.image.BufferedImage; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; + +import org.openqa.selenium.OutputType; +import org.openqa.selenium.WebDriver; +import org.openqa.selenium.chromium.HasCdp; + +import pazone.ashot.coordinates.Coords; +import pazone.ashot.util.ImageTool; + +/** + * Gets a screenshot using + * + * capture screenshot function provided by Chrome DevTools protocol. {@link WebDriver} instance provided + * to the class methods must be an instance of {@link HasCdp} and support Chrome DevTools protocol. + */ +public class CdpShootingStrategy implements ShootingStrategy { + + private static final long serialVersionUID = -4371668803381640029L; + + @Override + public BufferedImage getScreenshot(WebDriver driver) { + return getScreenshot(driver, Set.of()); + } + + @Override + public BufferedImage getScreenshot(WebDriver driver, Set coords) { + if (!HasCdp.class.isAssignableFrom(driver.getClass())) { + throw new IllegalArgumentException("WebDriver instance must support Chrome DevTools protocol"); + } + + Map args = new HashMap<>(); + args.put("captureBeyondViewport", true); + + if (!coords.isEmpty()) { + Coords elementCoords = coords.iterator().next(); + args.put("clip", Map.of( + "x", elementCoords.x, + "y", elementCoords.y, + "width", elementCoords.width, + "height", elementCoords.height, + "scale", 1) + ); + } + + Map results = ((HasCdp) driver).executeCdpCommand("Page.captureScreenshot", args); + String base64 = (String) results.get("data"); + byte[] bytes = OutputType.BYTES.convertFromBase64Png(base64); + + try { + return ImageTool.toBufferedImage(bytes); + } catch (IOException thrown) { + throw new UncheckedIOException(thrown); + } + } + + @Override + public Set prepareCoords(Set coordsSet) { + return coordsSet; + } +} diff --git a/src/test/java/pazone/ashot/CdpShootingStrategyTest.java b/src/test/java/pazone/ashot/CdpShootingStrategyTest.java new file mode 100644 index 0000000..c480517 --- /dev/null +++ b/src/test/java/pazone/ashot/CdpShootingStrategyTest.java @@ -0,0 +1,66 @@ +package pazone.ashot; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.awt.image.BufferedImage; +import java.io.IOException; +import java.util.Base64; +import java.util.Map; +import java.util.Set; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.openqa.selenium.WebDriver; +import org.openqa.selenium.chromium.HasCdp; + +import pazone.ashot.coordinates.Coords; +import pazone.ashot.util.ImageTool; +import pazone.ashot.util.TestImageUtils; + +@ExtendWith(MockitoExtension.class) +class CdpShootingStrategyTest { + private final ShootingStrategy strategy = new CdpShootingStrategy(); + + @Mock(extraInterfaces = HasCdp.class) + private WebDriver webDriver; + + @Test + void testPageScreenshot() throws IOException { + BufferedImage expected = TestImageUtils.IMAGE_A_SMALL; + String base = Base64.getEncoder().encodeToString(ImageTool.toByteArray(expected)); + + when(((HasCdp) webDriver).executeCdpCommand("Page.captureScreenshot", Map.of("captureBeyondViewport", true))) + .thenReturn(Map.of("data", base)); + + BufferedImage actual = strategy.getScreenshot(webDriver); + TestImageUtils.assertImageEquals(actual, expected); + } + + @Test + void testElementScreenshot() throws IOException { + BufferedImage expected = TestImageUtils.IMAGE_A_SMALL; + String base = Base64.getEncoder().encodeToString(ImageTool.toByteArray(expected)); + Coords coords = new Coords(1, 2, 3, 4); + + when(((HasCdp) webDriver).executeCdpCommand("Page.captureScreenshot", Map.of("captureBeyondViewport", true, + "clip", + Map.of("x", coords.x, "y", coords.y, "width", coords.width, "height", coords.height, "scale", 1)))) + .thenReturn(Map.of("data", base)); + + BufferedImage actual = strategy.getScreenshot(webDriver, Set.of(coords)); + TestImageUtils.assertImageEquals(actual, expected); + } + + @Test + void testUnsupportedCdp() { + WebDriver driver = mock(WebDriver.class); + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, + () -> strategy.getScreenshot(driver)); + assertEquals("WebDriver instance must support Chrome DevTools protocol", thrown.getMessage()); + } +}