From 9afe5d37013f3a695d209d17aba58b3377764573 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Corentin=20S=C3=A9chet?= <csechet@entrouvert.com>
Date: Mon, 11 Apr 2022 14:49:57 +0200
Subject: [PATCH] screenshot: only give screen width as argument, as screenshot
 take full pages

---
 frontools/cli.py        |  6 +++---
 frontools/screenshot.py | 24 +++++++++++-------------
 frontools/sources.py    | 18 +++++++++++-------
 3 files changed, 25 insertions(+), 23 deletions(-)

diff --git a/frontools/cli.py b/frontools/cli.py
index 1e489dd..9157064 100644
--- a/frontools/cli.py
+++ b/frontools/cli.py
@@ -97,17 +97,17 @@ async def css_diff_cli(config: Config, right_source: str) -> None:
 @main.command(name="screenshot-diff")
 @argument("source", type=str)
 @option("-o", "--output-directory", type=PathArgument(), default=None)
-@option("-r", "--resolution", type=str, default=None)
+@option("-w", "--screen-width", type=int, default=None)
 @_async_command
 @pass_obj
 async def screenshot_diff_cli(
     config: Config,
     source: str,
     output_directory: Optional[str],
-    resolution: Optional[str],
+    screen_width: Optional[int],
 ) -> None:
     """Generate screenshot diffs"""
-    await screenshot_diff(config, source, output_directory, resolution=resolution)
+    await screenshot_diff(config, source, output_directory, screen_width=screen_width)
 
 
 if __name__ == "__main__":
diff --git a/frontools/screenshot.py b/frontools/screenshot.py
index 033a941..8d7648e 100644
--- a/frontools/screenshot.py
+++ b/frontools/screenshot.py
@@ -4,6 +4,7 @@ from tempfile import NamedTemporaryFile
 from typing import Optional
 
 from PIL import Image, ImageChops
+from playwright.async_api import Error as PlaywrightError
 
 from frontools.config import Config
 from frontools.sources import Browser
@@ -18,7 +19,7 @@ async def screenshot_diff(
     config: Config,
     right_source_name: str,
     output_directory: Optional[str],
-    resolution: Optional[str] = None,
+    screen_width: Optional[int] = None,
 ) -> None:
     """Compare pages with or without local css"""
 
@@ -31,17 +32,8 @@ async def screenshot_diff(
     left_source = config.default_source
     right_source = config.get_source(right_source_name)
 
-    width: Optional[int] = None
-    height: Optional[int] = None
-    if resolution is not None:
-        [width_str, height_str] = resolution.split("x")
-        width = int(width_str)
-        height = int(height_str)
-
-    async with left_source.get_browser(width=width, height=height) as left_browser:
-        async with right_source.get_browser(
-            width=width, height=height
-        ) as right_browser:
+    async with left_source.get_browser(width=screen_width) as left_browser:
+        async with right_source.get_browser(width=screen_width) as right_browser:
             urls = [
                 (site_name, url)
                 for (site_name, site) in config.sites
@@ -103,4 +95,10 @@ async def _diff_url(
 
 async def _screenshot_url(browser: Browser, url: str) -> bytes:
     async with browser.load_page(url) as page:
-        return await page.screenshot(full_page=True)
+        try:
+            return await page.screenshot(full_page=True)
+        except PlaywrightError:
+            pass
+
+        # Exception raisen when taking a screenshot of a to large page, retry without full_page
+        return await page.screenshot()
diff --git a/frontools/sources.py b/frontools/sources.py
index 1b34ca9..eee4db0 100644
--- a/frontools/sources.py
+++ b/frontools/sources.py
@@ -5,7 +5,7 @@ from re import Pattern
 from re import compile as re_compile
 from typing import AsyncGenerator, AsyncIterable, Optional, cast
 
-from aiohttp import ClientSession
+from aiohttp import ClientSession, ClientConnectionError
 from bs4 import BeautifulSoup
 from playwright.async_api import (
     BrowserContext,
@@ -30,7 +30,7 @@ class Browser:
         page = await self._browser_context.new_page()
         await page.route("*", self._source.route)
         await page.goto(url)
-        await page.wait_for_load_state("networkidle")
+        await page.wait_for_load_state("networkidle", timeout=1000 * 60 * 2)
         yield page
         await page.close()
 
@@ -83,11 +83,15 @@ class CachedSource(Source):
         """Get a page content from the local or remote cache."""
         return await self._cache.get(url, self._load_url)
 
-    @staticmethod
-    async def _load_url(url: str) -> bytes:
-        async with ClientSession() as session:
-            async with session.get(url) as response:
-                return await response.content.read()
+    async def _load_url(self, url: str) -> bytes:
+        try:
+            async with ClientSession() as session:
+                async with session.get(url) as response:
+                    return await response.content.read()
+        except ClientConnectionError as ex:
+            self._error_summary.add_error(f'error while loading {url} : {ex}')
+
+        return b''
 
 
 class OverrideSource(Source):