Commit 57386964 authored by Watchtek's avatar Watchtek
Browse files

#6 add flag - take screenshot onlt on first and last session

parent f9ee9e00
Loading
Loading
Loading
Loading
+52 −7
Original line number Diff line number Diff line
@@ -103,6 +103,22 @@ class MultipleTest(unittest.TestCase):
        os.makedirs(self.screenshots_folder, exist_ok=True)
        self.screenshot_base_dir = self.screenshots_folder

        # Debug screenshot toggle:
        # - Environment variable SCREENSHOT_ALL=1 / true / y will enable all screenshots (non-interactive)
        # - Otherwise ask interactively (default No)
        env_flag = os.getenv("SCREENSHOT_ALL", "").lower()
        if env_flag in ("1", "true", "y", "yes"):
            self.screenshot_all = True
            print("[INFO] SCREENSHOT_ALL environment variable enabled: saving screenshots for ALL drivers (debug mode).")
        else:
            print("\n6. [선택] 모든 세션에서 스크린샷을 저장하시겠습니까? (y/N): ")
            ans = input().strip().lower()
            self.screenshot_all = ans in ("y", "yes")
            if self.screenshot_all:
                print("[INFO] Debug mode ON: saving screenshots for ALL drivers.")
            else:
                print("[INFO] Default mode: saving screenshots only for first and last drivers.")

    def check_http_status(self, url: str, max_retry: int = 3, ignore_cert: bool = True) -> bool:
        """
        Checks the HTTP status of the given URL.
@@ -158,20 +174,49 @@ class MultipleTest(unittest.TestCase):
        self.accept_next_alert = True

    def _take_screenshot(self, driver, action_name):
        """Saves a screenshot with a unique, descriptive filename."""
        """Saves screenshots according to the screenshot_all flag.

        - If self.screenshot_all is True -> save for every driver.
        - Otherwise only save for driver.instance_number == 1 or == self.max_connections.
        """
        try:
            inst = getattr(driver, "instance_number", None)
            if inst is None:
                # No instance number set; skip silently
                return

            # If debug flag set -> save for everyone
            if getattr(self, "screenshot_all", False):
                do_save = True
            else:
                last_inst = getattr(self, "max_connections", None)
                # Default behavior: only save for instance 1 and planned last instance
                do_save = (inst == 1) or (last_inst is not None and inst == last_inst)

            if not do_save:
                return

            # proceed to save screenshot for allowed instances
            time.sleep(0.2)
            session_folder = os.path.join(self.screenshot_base_dir, f"session_{driver.instance_number}")
            session_folder = os.path.join(self.screenshot_base_dir, f"session_{inst}")
            os.makedirs(session_folder, exist_ok=True)
            counter = getattr(driver, 'screenshot_counter', 0) + 1
            setattr(driver, 'screenshot_counter', counter)

            counter = getattr(driver, "screenshot_counter", 0) + 1
            setattr(driver, "screenshot_counter", counter)
            counter_str = str(counter).zfill(4)
            filename = os.path.join(session_folder, f"{counter_str}_{action_name}.png")
            time.sleep(3) # Set common wait time before save screenshot(except custom dashboard)

            # allow page to settle (your previous behavior)
            time.sleep(3)
            driver.save_screenshot(filename)
            # optionally show debug
            # print(f"Screenshot saved: {filename}")
        except (FileNotFoundError, PermissionError):
            print(f"Error saving screenshot for session {driver.instance_number}")

        except (FileNotFoundError, PermissionError) as e:
            print(f"Error saving screenshot for session {getattr(driver, 'instance_number', 'unknown')}: {e}")
        except Exception as e:
            print(f"Unexpected error in _take_screenshot for session {getattr(driver, 'instance_number', 'unknown')}: {e}")


    def _interact_with_time_filters(self, driver):
        """