Skip to content

Commit d4ddbe8

Browse files
glibasglibas
committed
Add conditional waiting support before taking screenshot
1 parent 44b73a9 commit d4ddbe8

File tree

2 files changed

+58
-50
lines changed

2 files changed

+58
-50
lines changed

‎src/main/java/com/assertthat/selenium_shutterbug/core/Shutterbug.java

+32-39
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,22 @@
99
import com.assertthat.selenium_shutterbug.utils.web.ScrollStrategy;
1010
import org.openqa.selenium.WebDriver;
1111
import org.openqa.selenium.WebElement;
12+
import org.openqa.selenium.support.ui.ExpectedCondition;
13+
14+
import java.util.function.Function;
1215

1316
/**
1417
* Created by Glib_Briia on 26/06/2016.
1518
*/
1619
public class Shutterbug {
1720

1821
private static final int DEFAULT_SCROLL_TIMEOUT = 100;
22+
private static Function<WebDriver,?> beforeShootCondition;
23+
private static int beforeShootTimeout;
24+
25+
private Shutterbug(){
26+
27+
}
1928

2029
/**
2130
* Make screen shot of the viewport only.
@@ -73,34 +82,30 @@ public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, in
7382
}
7483

7584
/**
76-
* To be used when screen shooting the page
77-
* and need to scroll while making screen shots, either vertically or
78-
* horizontally or both directions (Chrome).
85+
* Wait for condition to be true before taking screenshot
7986
*
80-
* @param driver WebDriver instance
81-
* @param scroll ScrollStrategy How you need to scroll
82-
* @param betweenScrollTimeout Timeout to wait after scrolling and before taking screenshot
83-
* @param afterScrollTimeout Timeout to wait after scrolling and before taking screenshot
84-
* @return PageSnapshot instance
87+
* @param cond condition
88+
* @param timeout timeout wait for condition
89+
* @return Shutterbug
8590
*/
86-
public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, int betweenScrollTimeout, int afterScrollTimeout) {
87-
return shootPage(driver, scroll, betweenScrollTimeout, true, afterScrollTimeout);
91+
public static Shutterbug wait(ExpectedCondition<?> cond, int timeout) {
92+
beforeShootCondition = cond;
93+
beforeShootTimeout = timeout;
94+
return null;
8895
}
8996

9097
/**
91-
* To be used when screen shooting the page
92-
* and need to scroll while making screen shots, either vertically or
93-
* horizontally or both directions (Chrome).
98+
* Wait for before taking screenshot
9499
*
95-
* @param driver WebDriver instance
96-
* @param scroll ScrollStrategy How you need to scroll
97-
* @param useDevicePixelRatio whether or not take into account device pixel ratio
98-
* @return PageSnapshot instance
100+
* @param timeout timeout wait for condition
101+
* @return Shutterbug
99102
*/
100-
public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, boolean useDevicePixelRatio) {
101-
return shootPage(driver, scroll, 0, useDevicePixelRatio);
103+
public static Shutterbug wait(int timeout) {
104+
beforeShootTimeout = timeout;
105+
return null;
102106
}
103107

108+
104109
/**
105110
* To be used when screen shooting the page
106111
* and need to scroll while making screen shots, either vertically or
@@ -109,11 +114,10 @@ public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, bo
109114
* @param driver WebDriver instance
110115
* @param scroll ScrollStrategy How you need to scroll
111116
* @param useDevicePixelRatio whether or not take into account device pixel ratio
112-
* @param afterScrollTimeout Timeout to wait after scrolling and before taking screenshot
113117
* @return PageSnapshot instance
114118
*/
115-
public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, boolean useDevicePixelRatio, int afterScrollTimeout) {
116-
return shootPage(driver, scroll, 0, useDevicePixelRatio, afterScrollTimeout);
119+
public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, boolean useDevicePixelRatio) {
120+
return shootPage(driver, scroll, 0, useDevicePixelRatio);
117121
}
118122

119123
/**
@@ -128,25 +132,14 @@ public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, bo
128132
* @return PageSnapshot instance
129133
*/
130134
public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, int betweenScrollTimeout, boolean useDevicePixelRatio) {
131-
return shootPage(driver, scroll, betweenScrollTimeout, useDevicePixelRatio, 0);
132-
}
133-
134-
/**
135-
* To be used when screen shooting the page
136-
* and need to scroll while making screen shots, either vertically or
137-
* horizontally or both directions (Chrome).
138-
*
139-
* @param driver WebDriver instance
140-
* @param scroll ScrollStrategy How you need to scroll
141-
* @param betweenScrollTimeout Timeout to wait between each scrolling operation
142-
* @param afterScrollTimeout Timeout to wait after scrolling and before taking screenshot
143-
* @param useDevicePixelRatio whether or not take into account device pixel ratio
144-
* @return PageSnapshot instance
145-
*/
146-
public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, int betweenScrollTimeout, boolean useDevicePixelRatio, int afterScrollTimeout) {
147135
Browser browser = new Browser(driver, useDevicePixelRatio);
148136
browser.setBetweenScrollTimeout(betweenScrollTimeout);
149-
browser.setAfterScrollTimeout(afterScrollTimeout);
137+
if(beforeShootCondition != null) {
138+
browser.setBeforeShootTimeout(beforeShootTimeout);
139+
browser.setBeforeShootCondition(beforeShootCondition);
140+
}else if(beforeShootTimeout!=0){
141+
browser.setBeforeShootTimeout(beforeShootTimeout);
142+
}
150143

151144
PageSnapshot pageScreenshot = new PageSnapshot(driver, browser.getDevicePixelRatio());
152145
switch (scroll) {

‎src/main/java/com/assertthat/selenium_shutterbug/utils/web/Browser.java

+26-11
Original file line numberDiff line numberDiff line change
@@ -8,20 +8,17 @@
88
import com.assertthat.selenium_shutterbug.utils.file.FileUtil;
99
import com.github.zafarkhaja.semver.Version;
1010
import com.google.common.collect.ImmutableMap;
11+
import org.openqa.selenium.*;
1112
import org.openqa.selenium.Dimension;
12-
import org.openqa.selenium.JavascriptExecutor;
13-
import org.openqa.selenium.OutputType;
1413
import org.openqa.selenium.Point;
15-
import org.openqa.selenium.TakesScreenshot;
16-
import org.openqa.selenium.WebDriver;
17-
import org.openqa.selenium.WebElement;
1814
import org.openqa.selenium.chrome.ChromeDriver;
1915
import org.openqa.selenium.firefox.FirefoxDriver;
2016
import org.openqa.selenium.remote.CommandInfo;
2117
import org.openqa.selenium.remote.HttpCommandExecutor;
2218
import org.openqa.selenium.remote.RemoteWebDriver;
2319
import org.openqa.selenium.remote.Response;
2420
import org.openqa.selenium.remote.http.HttpMethod;
21+
import org.openqa.selenium.support.ui.FluentWait;
2522

2623
import javax.imageio.ImageIO;
2724
import java.awt.*;
@@ -32,8 +29,10 @@
3229
import java.io.InputStream;
3330
import java.lang.reflect.InvocationTargetException;
3431
import java.lang.reflect.Method;
32+
import java.time.Duration;
3533
import java.util.ArrayList;
3634
import java.util.Map;
35+
import java.util.function.Function;
3736

3837
/**
3938
* Created by Glib_Briia on 17/06/2016.
@@ -59,8 +58,9 @@ public class Browser {
5958
private int viewportWidth = -1;
6059
private int viewportHeight = -1;
6160
private int betweenScrollTimeout;
61+
private Function<WebDriver,?> beforeShootCondition;
62+
private int beforeShootTimeout;
6263

63-
private int afterScrollTimeout;
6464
private Double devicePixelRatio = 1.0;
6565

6666
public Browser(WebDriver driver, boolean useDevicePixelRatio) {
@@ -83,16 +83,31 @@ public static void wait(int ms) {
8383
}
8484
}
8585

86+
public void wait(Function<WebDriver,?> condition, int timeout) {
87+
if(condition!=null) {
88+
new FluentWait<>(driver)
89+
.withTimeout(Duration.ofSeconds(timeout))
90+
.ignoring(StaleElementReferenceException.class, NoSuchMethodException.class)
91+
.until(condition);
92+
}else if(timeout>0) {
93+
wait(timeout);
94+
}
95+
}
96+
8697
public void setBetweenScrollTimeout(int betweenScrollTimeout) {
8798
this.betweenScrollTimeout = betweenScrollTimeout;
8899
}
89100

90-
public void setAfterScrollTimeout(int afterScrollTimeout) {
91-
this.afterScrollTimeout = afterScrollTimeout;
101+
public void setBeforeShootTimeout(int beforeShootTimeout) {
102+
this.beforeShootTimeout = beforeShootTimeout;
92103
}
93104

105+
public void setBeforeShootCondition(Function<WebDriver,?> beforeShootCondition) {
106+
this.beforeShootCondition = beforeShootCondition;
107+
}
94108

95109
public BufferedImage takeScreenshot() {
110+
wait(beforeShootCondition,beforeShootTimeout);
96111
File srcFile = ((TakesScreenshot) this.getUnderlyingDriver()).getScreenshotAs(OutputType.FILE);
97112
try {
98113
return ImageIO.read(srcFile);
@@ -157,7 +172,7 @@ public BufferedImage takeScreenshotEntirePageDefault() {
157172
final int scrollBarMaxWidth = 40; // this is probably too high, but better to be safe than sorry
158173

159174
if (_viewportWidth < _docWidth || (_viewportHeight < _docHeight && _viewportWidth - scrollBarMaxWidth < _docWidth))
160-
_viewportHeight -= scrollBarMaxWidth; // some space for a scrollbar
175+
_viewportHeight -= scrollBarMaxWidth; // some space for a scrollbar
161176
if (_viewportHeight < _docHeight)
162177
_viewportWidth -= scrollBarMaxWidth; // some space for a scrollbar
163178

@@ -194,7 +209,7 @@ public BufferedImage takeScreenshotEntirePageUsingChromeCommand() {
194209
}
195210
Object metrics = this.evaluate(FileUtil.getJsScript(ALL_METRICS));
196211
this.sendCommand("Emulation.setDeviceMetricsOverride", metrics);
197-
wait(afterScrollTimeout);
212+
wait(beforeShootCondition,beforeShootTimeout);
198213
Object result = this.sendCommand("Page.captureScreenshot", ImmutableMap.of("format", "png", "fromSurface", true));
199214
this.sendCommand("Emulation.clearDeviceMetricsOverride", ImmutableMap.of());
200215
return decodeBase64EncodedPng((String) ((Map<String, ?>) result).get("data"));
@@ -268,7 +283,7 @@ public void scrollToElement(WebElement element) {
268283
public void scrollToElementVerticalCentered(WebElement element) {
269284
executeJsScript(SCROLL_INTO_VIEW_VERTICAL_CENTERED_JS, element);
270285
}
271-
286+
272287
public void scrollTo(int x, int y) {
273288
executeJsScript(SCROLL_TO_JS, x / devicePixelRatio, y / devicePixelRatio);
274289
}

0 commit comments

Comments
 (0)