Coverage for src / ocarina / infra / screenshotter.py: 0.00%

1"""Generic screenshot capture utility for test automation. 

2 

3This module provides a reusable Screenshotter class that works with any 

4screenshot-capable driver through dependency injection. 

5 

6The library is framework-agnostic and can be adapted to: 

7- Selenium WebDriver 

8- Playwright 

9- Puppeteer 

10- Any custom screenshot driver 

11 

12Key features: 

13- Protocol-based (driver must have save_screenshot) 

14- Optional full-page screenshot via injection 

15- Generic over actual driver type 

16- Configurable via ScreenshotterConfig 

17- Thread-safe operations 

18- Burst mode support 

19- Collision-free filename generation 

20 

21Usage: 

22 >>> # Create config for your project 

23 >>> config = ScreenshotterConfig[WebDriver]( 

24 ... output_dir=Path(".screenshots"), 

25 ... file_ext=".png", 

26 ... health_check=selenium_health_check, 

27 ... save_full_page=selenium_save_full_page, # Optional 

28 ... ) 

29 >>> screenshotter = Screenshotter(driver, logger, config) 

30 >>> screenshotter.take_screenshot(prefix="test") 

31""" 

32 

33import time 

34import uuid 

35from collections.abc import Callable 

36from dataclasses import dataclass 

37from pathlib import Path 

38from threading import Lock 

39from typing import TYPE_CHECKING, Final, Protocol, TypeVar 

40 

41if TYPE_CHECKING: 

42 from ocarina.ports.ilogger import ILogger 

43 

44 

45class ScreenshotDriver(Protocol): 

46 """Protocol for drivers with screenshot capability. 

47 

48 Drivers must implement save_screenshot() method. 

49 Full-page screenshot is optional and injected via config. 

50 """ 

51 

52 def save_screenshot(self, path: str) -> bool: 

53 """Save standard viewport screenshot. 

54 

55 Args: 

56 path: File path where screenshot should be saved. 

57 

58 Returns: 

59 True if successful, False otherwise. 

60 

61 """ 

62 ... 

63 

64 

65TDriver = TypeVar("TDriver", bound=ScreenshotDriver) 

66"""Type variable for driver types implementing ScreenshotDriver protocol. 

67 

68Bounded to ScreenshotDriver to ensure save_screenshot() exists, 

69while preserving the concrete type for health checks and full-page screenshots. 

70""" 

71 

72HealthCheck = Callable[[TDriver], None] 

73"""Health check function type. 

74 

75Generic over actual driver type. Allows health checks to use 

76driver-specific methods from the concrete driver type. 

77 

78Example: 

79 >>> # Selenium 

80 ... def selenium_health_check(driver: WebDriver) -> None: 

81 ... _ = driver.title # WebDriver-specific 

82 ... _ = driver.current_url # Any WebDriver method 

83""" 

84 

85SaveFullPageScreenshot = Callable[[TDriver, str], bool] 

86"""Full-page screenshot function type. 

87 

88Optional function injected via config. Allows framework-specific 

89implementations without polluting the core library. 

90 

91Example: 

92 >>> # Selenium Firefox 

93 ... def selenium_save_full_page(driver: WebDriver, path: str) -> bool: 

94 ... if hasattr(driver, "save_full_page_screenshot"): 

95 ... return cast(FirefoxWebDriver, driver).save_full_page_screenshot(path) 

96 ... return False 

97""" 

98 

99SCREENSHOT_SUCCESS_PREFIX: Final[str] = "Screenshot: " 

100 

101 

102@dataclass(frozen=True) 

103class ScreenshotterConfig[TDriver: ScreenshotDriver]: 

104 """Configuration for Screenshotter behavior. 

105 

106 Allows each project to customize screenshot behavior without modifying 

107 the library code. 

108 

109 Type Parameters: 

110 TDriver: Driver type implementing ScreenshotDriver protocol. 

111 

112 Example: 

113 >>> # Selenium project 

114 ... config = ScreenshotterConfig[WebDriver]( 

115 ... output_dir=Path.cwd() / ".screenshots", 

116 ... file_ext=".png", 

117 ... health_check=selenium_health_check, 

118 ... save_full_page=selenium_save_full_page, 

119 ... ) 

120 

121 Attributes: 

122 output_dir: Directory where screenshots are saved. 

123 file_ext: File extension for screenshots (e.g., ".png", ".jpg"). 

124 health_check: Optional function to check driver health before screenshot. 

125 save_full_page: Optional function for full-page screenshots. 

126 If None, only standard screenshots are taken. 

127 default_burst_delay: Default delay between burst shots in seconds. 

128 max_filename_retries: Max attempts to generate unique filename. 

129 uuid_length: Length of UUID suffix in filename. 

130 

131 """ 

132 

133 output_dir: Path 

134 file_ext: str = ".png" 

135 health_check: HealthCheck[TDriver] | None = None 

136 save_full_page: SaveFullPageScreenshot[TDriver] | None = None 

137 default_burst_delay: float = 0.5 

138 max_filename_retries: int = 500 

139 uuid_length: int = 8 

140 

141 

142_SCREENSHOTTER_LOCK = Lock() 

143"""Global lock ensuring thread-safe screenshot operations. 

144 

145Prevents race conditions in filename generation and file writes 

146across multiple threads. 

147""" 

148 

149 

150class Screenshotter[TDriver: ScreenshotDriver]: 

151 """Thread-safe screenshot capture utility. 

152 

153 Generic implementation that works with any driver implementing 

154 ScreenshotDriver protocol. Configuration injected via ScreenshotterConfig. 

155 

156 Type Parameters: 

157 TDriver: Driver type implementing ScreenshotDriver protocol. 

158 

159 Attributes: 

160 _driver: Screenshot-capable driver. 

161 _logger: Logger for recording operations. 

162 _config: Configuration controlling behavior. 

163 _output_dir: Directory for screenshots (from config). 

164 

165 Example: 

166 >>> config = ScreenshotterConfig[WebDriver](output_dir=Path(".screenshots")) 

167 >>> screenshotter = Screenshotter(driver, logger, config) 

168 >>> screenshotter.take_screenshot(prefix="test") 

169 

170 """ 

171 

172 def __init__( 

173 self, 

174 driver: TDriver, 

175 logger: ILogger, 

176 config: ScreenshotterConfig[TDriver], 

177 ) -> None: 

178 """Initialize screenshotter with driver, logger, and config. 

179 

180 Creates output directory if it doesn't exist. 

181 

182 Args: 

183 driver: Driver implementing ScreenshotDriver protocol. 

184 logger: Logger implementing ILogger interface. 

185 config: Configuration for screenshotter behavior. 

186 

187 """ 

188 self._driver = driver 

189 self._logger = logger 

190 self._config = config 

191 self._output_dir = config.output_dir 

192 self._output_dir.mkdir(parents=True, exist_ok=True) 

193 

194 def _generate_unique_file_path(self, *, prefix: str, counter: int) -> Path | None: 

195 """Generate unique screenshot filename with collision avoidance. 

196 

197 Args: 

198 prefix: Optional prefix for filename. 

199 counter: Shot number in burst (-1 for single shot). 

200 

201 Returns: 

202 Path if unique filename generated, None if all retries exhausted. 

203 

204 """ 

205 retries = self._config.max_filename_retries 

206 uuid_length = self._config.uuid_length 

207 burst = counter != -1 

208 

209 for _ in range(retries): 

210 unique_id = uuid.uuid4().hex[:uuid_length] 

211 file_name = f"{prefix}_{unique_id}" if prefix else f"{unique_id}" 

212 base_path = self._output_dir / file_name 

213 

214 normalized_file_path = ( 

215 f"{base_path}_{counter}{self._config.file_ext}" 

216 if burst 

217 else f"{base_path}{self._config.file_ext}" 

218 ) 

219 

220 if not Path(normalized_file_path).exists(): 

221 return Path(normalized_file_path) 

222 

223 return None 

224 

225 def take_screenshot( 

226 self, 

227 *, 

228 prefix: str = "", 

229 shots: int | None = None, 

230 burst_delay: float | None = None, 

231 ) -> None: 

232 """Capture one or more screenshots with automatic file naming. 

233 

234 Thread-safe operation with optional health check and burst mode support. 

235 

236 Example: 

237 >>> # Single shot 

238 ... screenshotter.take_screenshot(prefix="login") 

239 

240 Example: 

241 >>> # Burst mode 

242 ... screenshotter.take_screenshot( 

243 ... prefix="animation", 

244 ... shots=5, 

245 ... burst_delay=0.2 

246 ... ) 

247 

248 Args: 

249 prefix: Optional prefix for screenshot filename. 

250 shots: Number of screenshots to take (default: 1). 

251 burst_delay: Delay between burst shots in seconds 

252 (default: from config). 

253 

254 """ 

255 dead_driver_msg = "Cannot take screenshot, driver died." 

256 

257 def _check_driver_health() -> Exception | None: 

258 """Run health check if configured.""" 

259 if self._config.health_check is None: 

260 return None 

261 

262 try: 

263 self._config.health_check(self._driver) 

264 except Exception as exc: # noqa: BLE001 

265 return exc 

266 return None 

267 

268 dead_driver_exc = _check_driver_health() 

269 

270 if dead_driver_exc: 

271 self._logger.exception(dead_driver_msg, exc=dead_driver_exc) 

272 return 

273 

274 if shots is None: 

275 shots = 1 

276 

277 if burst_delay is None: 

278 burst_delay = self._config.default_burst_delay 

279 

280 burst = shots > 1 

281 

282 with _SCREENSHOTTER_LOCK: 

283 for i in range(1, shots + 1): 

284 if burst and i > 1: 

285 time.sleep(burst_delay) 

286 

287 normalized_file_path = self._generate_unique_file_path( 

288 prefix=prefix, counter=i if burst else -1 

289 ) 

290 

291 if normalized_file_path is None: 

292 msg = ( 

293 "FAILED TO TAKE SCREENSHOT!" 

294 " " 

295 f"(Can't generate unique file path for prefix '{prefix}')" 

296 ) 

297 self._logger.error(msg) 

298 continue 

299 

300 success = False 

301 if self._config.save_full_page is not None: 

302 success = self._config.save_full_page( 

303 self._driver, str(normalized_file_path) 

304 ) 

305 

306 if not success: 

307 success = self._driver.save_screenshot(str(normalized_file_path)) 

308 

309 if success: 

310 msg = f"{SCREENSHOT_SUCCESS_PREFIX}{normalized_file_path}" 

311 self._logger.info(msg) 

312 else: 

313 dead_driver_exc = _check_driver_health() 

314 if dead_driver_exc: 

315 self._logger.exception(dead_driver_msg, exc=dead_driver_exc) 

316 return 

317 

318 msg = f"FAILED TO TAKE SCREENSHOT! ({normalized_file_path})" 

319 self._logger.error(msg)