Remove --bundle option from screenshot.py and add a full page test

palewire · Nov 20, 2022 · bba0676 · bba0676
1 parent 6c0142d
commit bba0676
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 24 deletions.
diff --git a/newshomepages/screenshot.py b/newshomepages/screenshot.py
@@ -24,47 +24,35 @@
     default=False,
     help="Screenshot the whole page",
 )
-@click.option(
-    "--bundle",
-    "is_bundle",
-    is_flag=True,
-    default=False,
-    help="The provided handle is a bundle",
-)
 def cli(
     handle: str,
     output_dir: str = "./",
     wait: str = "5000",
     width: str = "1300",
     height: str = "1600",
     full_page: bool = False,
-    is_bundle: bool = False,
 ):
     """Screenshot the provided homepage."""
     # Set the output path
     output_path = Path(output_dir)
     output_path.mkdir(parents=True, exist_ok=True)
 
-    if is_bundle:
-        site_list = utils.get_sites_in_bundle(handle)
-    else:
-        site_list = [utils.get_site(handle)]
+    # Get the site
+    site = utils.get_site(handle)
 
     # Open the browser
     with sync_playwright() as playwright:
         # We'll load it with an extension
         context = utils._load_persistent_context(playwright, int(width), int(height))
 
-        # Loop through all the sites
-        for site in site_list:
-            # Screenshot them one by one
-            _screenshot(
-                context,
-                site,
-                output_path,
-                wait=int(site["wait"] or wait),
-                full_page=full_page,
-            )
+        # Screenshot the site
+        _screenshot(
+            context,
+            site,
+            output_path,
+            wait=int(site["wait"] or wait),
+            full_page=full_page,
+        )
 
         # Close it out
         context.close()

diff --git a/tests/test_screenshot.py b/tests/test_screenshot.py
@@ -3,8 +3,15 @@
 from newshomepages import screenshot
 
 
-def test_single(tmp_path):
-    """Test a single screenshot."""
+def test_cropped(tmp_path):
+    """Test a cropped screenshot."""
     runner = CliRunner()
     result = runner.invoke(screenshot.cli, ["latimes", "-o", tmp_path])
     assert result.exit_code == 0
+
+
+def test_full(tmp_path):
+    """Test a full-page screenshot."""
+    runner = CliRunner()
+    result = runner.invoke(screenshot.cli, ["latimes", "-o", tmp_path, "--full-page"])
+    assert result.exit_code == 0