Add "--no-screenshot" option to skip saving screenshots

Author: mdmintz

seleniumbase/core/log_helper.py

@@ -12,7 +12,15 @@
 def log_screenshot(test_logpath, driver, screenshot=None, get=False):
     screenshot_name = settings.SCREENSHOT_NAME
     screenshot_path = os.path.join(test_logpath, screenshot_name)
+    screenshot_skipped = constants.Warnings.SCREENSHOT_SKIPPED
     screenshot_warning = constants.Warnings.SCREENSHOT_UNDEFINED
+    if (
+        (hasattr(sb_config, "no_screenshot") and sb_config.no_screenshot)
+        or screenshot == screenshot_skipped
+    ):
+        if get:
+            return screenshot
+        return
     try:
         if not screenshot:
             element = driver.find_element_by_tag_name("body")

seleniumbase/fixtures/base_case.py

@@ -12925,6 +12925,7 @@ def setUp(self, masterqa_mode=False):
                 settings.HEADLESS_START_HEIGHT = height
             self.maximize_option = sb_config.maximize_option
             self.save_screenshot_after_test = sb_config.save_screenshot
+            self.no_screenshot_after_test = sb_config.no_screenshot
             self.visual_baseline = sb_config.visual_baseline
             self.timeout_multiplier = sb_config.timeout_multiplier
             self.pytest_html_report = sb_config.pytest_html_report
@@ -13226,38 +13227,73 @@ def setUp(self, masterqa_mode=False):
     def __set_last_page_screenshot(self):
         """self.__last_page_screenshot is only for pytest html report logs.
         self.__last_page_screenshot_png is for all screenshot log files."""
-        if not self.__last_page_screenshot and (
-            not self.__last_page_screenshot_png
+        SCREENSHOT_SKIPPED = constants.Warnings.SCREENSHOT_SKIPPED
+        SCREENSHOT_UNDEFINED = constants.Warnings.SCREENSHOT_UNDEFINED
+        if (
+            hasattr(self, "no_screenshot_after_test")
+            and self.no_screenshot_after_test
+        ):
+            from seleniumbase.core import encoded_images
+
+            NO_SCREENSHOT = encoded_images.get_no_screenshot_png()
+            self.__last_page_screenshot = NO_SCREENSHOT
+            self.__last_page_screenshot_png = SCREENSHOT_SKIPPED
+            sb_config._last_page_screenshot_png = NO_SCREENSHOT
+            return
+        element = None
+        if (
+            not self.__last_page_screenshot
+            and not self.__last_page_screenshot_png
         ):
             try:
                 element = self.driver.find_element(by="tag name", value="body")
-                if self.is_pytest and self.report_on:
-                    self.__last_page_screenshot_png = (
-                        self.driver.get_screenshot_as_png()
-                    )
+                try:
                     self.__last_page_screenshot = element.screenshot_as_base64
-                else:
-                    self.__last_page_screenshot_png = element.screenshot_as_png
-            except Exception:
-                if not self.__last_page_screenshot:
-                    if self.is_pytest and self.report_on:
-                        try:
-                            self.__last_page_screenshot = (
-                                self.driver.get_screenshot_as_base64()
-                            )
-                        except Exception:
-                            self.__last_page_screenshot = (
-                                constants.Warnings.SCREENSHOT_UNDEFINED
-                            )
-                if not self.__last_page_screenshot_png:
+                except Exception:
                     try:
-                        self.__last_page_screenshot_png = (
-                            self.driver.get_screenshot_as_png()
+                        self.__last_page_screenshot = (
+                            self.driver.get_screenshot_as_base64()
                         )
                     except Exception:
+                        pass
+            except Exception:
+                pass
+            if not self.__last_page_screenshot:
+                self.__last_page_screenshot = SCREENSHOT_UNDEFINED
+                self.__last_page_screenshot_png = SCREENSHOT_UNDEFINED
+                if element:
+                    try:
                         self.__last_page_screenshot_png = (
-                            constants.Warnings.SCREENSHOT_UNDEFINED
+                            element.screenshot_as_png
                         )
+                    except Exception:
+                        try:
+                            self.__last_page_screenshot_png = (
+                                self.driver.get_screenshot_as_png()
+                            )
+                        except Exception:
+                            pass
+            else:
+                import base64
+
+                try:
+                    self.__last_page_screenshot_png = (
+                        base64.b64decode(self.__last_page_screenshot)
+                    )
+                except Exception:
+                    if element:
+                        try:
+                            self.__last_page_screenshot_png = (
+                                element.screenshot_as_png
+                            )
+                        except Exception:
+                            try:
+                                self.__last_page_screenshot_png = (
+                                    self.driver.get_screenshot_as_png()
+                                )
+                            except Exception:
+                                pass
+        sb_config._last_page_screenshot_png = self.__last_page_screenshot_png
 
     def __set_last_page_url(self):
         if not self.__last_page_url:

seleniumbase/fixtures/constants.py

@@ -142,6 +142,7 @@ class Values:
 
 
 class Warnings:
+    SCREENSHOT_SKIPPED = "Skipping screenshot!"
     SCREENSHOT_UNDEFINED = "Unable to get screenshot!"
     PAGE_SOURCE_UNDEFINED = "Unable to get page source!"
     INVALID_RUN_COMMAND = "INVALID RUN COMMAND!"

seleniumbase/plugins/pytest_plugin.py

@@ -106,6 +106,7 @@ def pytest_addoption(parser):
     --window-size=WIDTH,HEIGHT  (Set the browser's starting window size.)
     --maximize  (Start tests with the browser window maximized.)
     --screenshot  (Save a screenshot at the end of each test.)
+    --no-screenshot  (No screenshots saved unless tests directly ask it.)
     --visual-baseline  (Set the visual baseline for Visual/Layout tests.)
     --wire  (Use selenium-wire's webdriver for replacing selenium webdriver.)
     --external-pdf  (Set Chromium "plugins.always_open_pdf_externally":True.)
@@ -1083,8 +1084,20 @@ def pytest_addoption(parser):
         action="store_true",
         dest="save_screenshot",
         default=False,
-        help="""Save a screenshot at the end of the test.
-                (Added to the "latest_logs/" folder.)""",
+        help="""Save a screenshot at the end of every test.
+                By default, this is only done for failures.
+                Will be saved in the "latest_logs/" folder.""",
+    )
+    parser.addoption(
+        "--no-screenshot",
+        "--no_screenshot",
+        "--ns",
+        action="store_true",
+        dest="no_screenshot",
+        default=False,
+        help="""No screenshots saved unless tests directly ask it.
+                This changes default behavior where screenshots are
+                saved for test failures and pytest-html reports.""",
     )
     parser.addoption(
         "--visual_baseline",
@@ -1435,6 +1448,7 @@ def pytest_configure(config):
     sb_config.window_size = config.getoption("window_size")
     sb_config.maximize_option = config.getoption("maximize_option")
     sb_config.save_screenshot = config.getoption("save_screenshot")
+    sb_config.no_screenshot = config.getoption("no_screenshot")
     sb_config.visual_baseline = config.getoption("visual_baseline")
     sb_config.use_wire = config.getoption("use_wire")
     sb_config.external_pdf = config.getoption("external_pdf")
@@ -1559,12 +1573,14 @@ def pytest_configure(config):
     if sb_config.dash_title:
         constants.Dashboard.TITLE = sb_config.dash_title.replace("_", " ")
 
+    if sb_config.save_screenshot and sb_config.no_screenshot:
+        sb_config.save_screenshot = False  # "no_screenshot" has priority
+
     if (
         sb_config._multithreaded
         and "--co" not in sys_argv
         and "--collect-only" not in sys_argv
     ):
-        from seleniumbase.core import log_helper
         from seleniumbase.core import download_helper
         from seleniumbase.core import proxy_helper
 

seleniumbase/plugins/selenium_plugin.py

@@ -770,8 +770,20 @@ def options(self, parser, env):
             action="store_true",
             dest="save_screenshot",
             default=False,
-            help="""Save a screenshot at the end of the test.
-                    (Added to the "latest_logs/" folder.)""",
+            help="""Save a screenshot at the end of every test.
+                    By default, this is only done for failures.
+                    Will be saved in the "latest_logs/" folder.""",
+        )
+        parser.add_option(
+            "--no-screenshot",
+            "--no_screenshot",
+            "--ns",
+            action="store_true",
+            dest="no_screenshot",
+            default=False,
+            help="""No screenshots saved unless tests directly ask it.
+                    This changes default behavior where screenshots are
+                    saved for test failures and pytest-html reports.""",
         )
         parser.add_option(
             "--visual_baseline",
@@ -948,14 +960,18 @@ def beforeTest(self, test):
         test.test._disable_beforeunload = self.options._disable_beforeunload
         test.test.window_size = self.options.window_size
         test.test.maximize_option = self.options.maximize_option
+        if self.options.save_screenshot and self.options.no_screenshot:
+            self.options.save_screenshot = False  # no_screenshot has priority
         test.test.save_screenshot_after_test = self.options.save_screenshot
+        test.test.no_screenshot_after_test = self.options.no_screenshot
         test.test.visual_baseline = self.options.visual_baseline
         test.test.use_wire = self.options.use_wire
         test.test.external_pdf = self.options.external_pdf
         test.test.timeout_multiplier = self.options.timeout_multiplier
         test.test.dashboard = False
         test.test._multithreaded = False
         test.test._reuse_session = False
+        sb_config.no_screenshot = test.test.no_screenshot_after_test
         if test.test.servername != "localhost":
             # Using Selenium Grid
             # (Set --server="127.0.0.1" for localhost Grid)