From 0cded6bd6414551af8bda15e62b8ec5713f13eb8 Mon Sep 17 00:00:00 2001 From: GoWithitRoger Date: Mon, 11 Aug 2025 11:04:29 -0500 Subject: [PATCH] refactor: move chromedriver cleanup into managed_webdriver_session context manager --- main.py | 135 ++++++++++++++++------------------------ tests/test_cleanup.py | 49 ++++++++++----- tests/test_main_flow.py | 25 ++------ 3 files changed, 93 insertions(+), 116 deletions(-) diff --git a/main.py b/main.py index 4807ef6..6e4404e 100644 --- a/main.py +++ b/main.py @@ -56,29 +56,6 @@ def set_chromedriver_pid(self, pid: int) -> None: self.last_chromedriver_pid = pid -def cleanup_old_processes(debug_logger: Optional["DebugLogger"] = None) -> None: - """Find and terminate any lingering chromedriver processes from previous runs. - - This is a safeguard against resource leaks from crashed/orphaned sessions. - """ - print("Running pre-emptive cleanup of old chromedriver processes...") - if debug_logger: - debug_logger.log("cleanup_old_processes: START") - try: - # Use pkill to find and kill processes by name. -f checks the full command line. - # The [c]haracter class is a trick to prevent pkill from finding its own process. - command = "pkill -f '[c]hromedriver'" - result = subprocess.run(command, shell=True, check=False) - if debug_logger: - rc = getattr(result, "returncode", "NA") - debug_logger.log(f"cleanup_old_processes: END (rc={rc})") - print("Cleanup complete.") - except Exception as e: - if debug_logger: - debug_logger.log(f"cleanup_old_processes: ERROR {e}") - print(f"Notice: Pre-emptive cleanup failed. This is non-critical. Error: {e}") - - def log_running_chromedriver_processes(debug_logger: DebugLogger) -> None: """Logs any active chromedriver processes using `ps`. @@ -102,52 +79,64 @@ def log_running_chromedriver_processes(debug_logger: DebugLogger) -> None: debug_logger.log(f"Error while checking chromedriver processes: {e}") +def cleanup_old_processes() -> None: + """Kill lingering chromedriver processes from prior runs (best-effort).""" + try: + subprocess.run("pkill -f '[c]hromedriver'", shell=True, check=False) + except Exception: + # Ignore any error; this is a best-effort cleanup. + pass + + @contextmanager def managed_webdriver_session(chrome_options: Options, debug_logger: DebugLogger): - """A self-cleaning context manager for the Selenium WebDriver. + """A self-contained, resilient context manager for Selenium WebDriver. - Sets up the driver and guarantees the chromedriver service process is - terminated on exit, aggressively if needed to avoid orphans. + It handles pre-emptive cleanup, robust initialization, and guaranteed teardown. """ - service: ChromeService = ChromeService() + debug_logger.log("managed_webdriver_session: START") + print("Running pre-emptive cleanup of old chromedriver processes...") + try: + subprocess.run("pkill -f '[c]hromedriver'", shell=True, check=False) + print("Cleanup complete.") + except Exception as e: + print(f"Notice: Pre-emptive cleanup failed. This is non-critical. Error: {e}") + driver: Optional[WebDriver] = None + service: Optional[ChromeService] = None print("Setting up WebDriver for gateway tests...") try: + service = ChromeService() driver = webdriver.Chrome(service=service, options=chrome_options) - # Track and log the chromedriver PID if available - try: - if service and service.process and service.process.pid: - debug_logger.set_chromedriver_pid(service.process.pid) - debug_logger.log( - f"WebDriver service started with PID: {service.process.pid}" - ) - except Exception as e: # pragma: no cover - best effort logging - debug_logger.log(f"Unable to retrieve WebDriver PID: {e}") + if service and service.process and service.process.pid: + debug_logger.set_chromedriver_pid(service.process.pid) + debug_logger.log(f"WebDriver service started with PID: {service.process.pid}") yield driver + except Exception as e: + print(f"CRITICAL: Failed to initialize WebDriver session. Error: {e}") + yield None finally: print("Shutting down WebDriver session...") if driver: debug_logger.log("WebDriver quit: START") try: - driver.quit() # Graceful shutdown of browser and driver + driver.quit() finally: debug_logger.log("WebDriver quit: END") - # Aggressive kill of chromedriver service process if service and getattr(service, "process", None): - pid = getattr(service.process, "pid", None) - if pid: - debug_logger.log( - f"Forcefully terminating chromedriver service (PID: {pid})..." - ) try: - service.process.kill() # Force kill to prevent orphans - service.process.wait(timeout=5) # Confirm termination - debug_logger.log("Service terminated successfully.") - except Exception as e: # pragma: no cover - defensive - debug_logger.log( - "Notice: Could not kill service process; it may have already exited. " + pid = getattr(service.process, "pid", None) + if pid: + print(f"Forcefully terminating chromedriver service (PID: {pid})...") + service.process.kill() + service.process.wait(timeout=5) + print("Service terminated successfully.") + except Exception as e: + print( + "Notice: Could not kill service process, it may have already exited. " f"Error: {e}" ) + debug_logger.log("managed_webdriver_session: END") # Verify process state after teardown log_running_chromedriver_processes(debug_logger) @@ -451,11 +440,8 @@ def run_ping_test_task(driver: WebDriver) -> Optional[GatewayPingResults]: return None return parse_gateway_ping_results(results_text) except Exception as e: - print(f"Error during gateway ping test: {e}") - error_time = datetime.now().strftime("%Y%m%d_%H%M%S") - screenshot_file = f"error_screenshot_gateway_ping_{error_time}.png" - driver.save_screenshot(screenshot_file) - print(f"Saved screenshot to {screenshot_file} for debugging.") + # Just log the error and return. Don't try to interact with a potentially dead driver. + print(f"An error occurred during the task: {e}") return None @@ -504,7 +490,8 @@ def run_speed_test_task(driver: WebDriver, access_code: str) -> Optional[SpeedRe break return results if results else None except Exception as e: - print(f"Error during gateway speed test: {e}") + # Just log the error and return. Don't try to interact with a potentially dead driver. + print(f"An error occurred during the task: {e}") return None @@ -713,14 +700,14 @@ def perform_checks() -> None: master_results: dict[str, str | float | int | None] = {} debug_log = DebugLogger(start_time=time.time()) debug_log.log("perform_checks: START") - # Pre-emptive scavenger to ensure no orphaned chromedriver processes linger - cleanup_old_processes(debug_log) print( f"\n[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] " f"Starting checks (Run #{run_counter})..." ) - # Check for any lingering chromedriver processes from previous runs + # Pre-emptive cleanup should run first + cleanup_old_processes() + # Optionally log existing chromedriver processes after cleanup in debug mode log_running_chromedriver_processes(debug_log) # --- Run Local Tests (No Browser Required) --- @@ -763,34 +750,20 @@ def perform_checks() -> None: chrome_options.add_argument("--no-sandbox") chrome_options.add_argument("--disable-dev-shm-usage") - try: - debug_log.log("Selenium setup: START") - with managed_webdriver_session(chrome_options, debug_log) as driver: - debug_log.log("Selenium setup: END") - driver.get(config.GATEWAY_URL) - time.sleep(2) + # The 'try...except' is removed, as the context manager handles its own errors. + with managed_webdriver_session(chrome_options, debug_log) as driver: + if driver: + # Safely run tasks if the driver was created successfully debug_log.log("run_ping_test_task: START") - gateway_ping_results = run_ping_test_task(driver) + master_results.update(run_ping_test_task(driver) or {}) debug_log.log("run_ping_test_task: END") - if gateway_ping_results: - master_results.update(gateway_ping_results) if should_run_gateway_speed_test: debug_log.log("run_speed_test_task: START") - gateway_speed_results = run_speed_test_task(driver, DEVICE_ACCESS_CODE) + master_results.update(run_speed_test_task(driver, DEVICE_ACCESS_CODE) or {}) debug_log.log("run_speed_test_task: END") - if gateway_speed_results: - master_results.update(gateway_speed_results) - except Exception as e: - print(f"A critical error occurred in the WebDriver session: {e}") - # Best-effort screenshot if driver existed during context - try: - if 'driver' in locals() and driver: - error_time = datetime.now().strftime("%Y%m%d_%H%M%S") - screenshot_file = f"error_screenshot_{error_time}.png" - driver.save_screenshot(screenshot_file) - print(f"Saved screenshot to {screenshot_file} for debugging.") - except Exception: - pass + else: + # This 'else' block will run if the context manager yields None + print("Skipping gateway tests because WebDriver session failed to start.") debug_log.log("perform_checks: END") log_results(master_results) diff --git a/tests/test_cleanup.py b/tests/test_cleanup.py index 7930a26..abf559b 100644 --- a/tests/test_cleanup.py +++ b/tests/test_cleanup.py @@ -1,6 +1,7 @@ import os import sys -from unittest.mock import patch +import time +from unittest.mock import MagicMock, patch # Ensure the main module can be imported sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) @@ -8,18 +9,38 @@ import main -def test_cleanup_old_processes_invokes_pkill(): - with patch("main.subprocess.run") as mock_run: - main.cleanup_old_processes() - mock_run.assert_called() - args, kwargs = mock_run.call_args - # Ensure pkill is targeted at chromedriver with -f and shell=True - assert "pkill -f '[c]hromedriver'" in args[0] - assert kwargs.get("shell") is True - assert kwargs.get("check") is False +def test_context_manager_runs_pkill_preemptively(): + debug_logger = main.DebugLogger(start_time=time.time()) + fake_service = MagicMock() + fake_service.process = MagicMock() + fake_driver = MagicMock() -def test_cleanup_old_processes_handles_exception(): - with patch("main.subprocess.run", side_effect=RuntimeError("boom")): - # Should not raise - main.cleanup_old_processes() + with ( + patch("main.subprocess.run") as mock_run, + patch.object(main, "ChromeService", return_value=fake_service), + patch.object(main.webdriver, "Chrome", return_value=fake_driver), + patch.object(main, "log_running_chromedriver_processes"), + ): + opts = main.Options() + with main.managed_webdriver_session(opts, debug_logger): + pass + # Pre-emptive cleanup should have been invoked + mock_run.assert_any_call("pkill -f '[c]hromedriver'", shell=True, check=False) + + +def test_context_manager_yields_none_on_setup_failure(): + debug_logger = main.DebugLogger(start_time=time.time()) + + with ( + patch("main.subprocess.run") as mock_run, + patch.object(main, "ChromeService") as mock_service, + patch.object(main.webdriver, "Chrome", side_effect=RuntimeError("boom")), + patch.object(main, "log_running_chromedriver_processes"), + ): + mock_service.return_value.process = MagicMock() + opts = main.Options() + with main.managed_webdriver_session(opts, debug_logger) as driver: + assert driver is None + # Ensure pkill was attempted even when setup fails + mock_run.assert_any_call("pkill -f '[c]hromedriver'", shell=True, check=False) diff --git a/tests/test_main_flow.py b/tests/test_main_flow.py index 0c82e8a..c420453 100644 --- a/tests/test_main_flow.py +++ b/tests/test_main_flow.py @@ -23,7 +23,7 @@ def mock_tasks(): patch("main.run_local_speed_test_task", return_value={}) as mock_local_speed, patch("main.run_ping_test_task", return_value={}) as mock_gateway_ping, patch("main.run_speed_test_task", return_value={}) as mock_gateway_speed, - patch("main.cleanup_old_processes") as mock_cleanup, + patch("main.subprocess.run") as mock_subproc_run, patch("main.log_results") as mock_log, patch("main.webdriver.Chrome", return_value=MagicMock()) as mock_chrome, patch("main.get_access_code", return_value="test-code") as mock_access_code, @@ -38,7 +38,7 @@ def mock_tasks(): "local_speed": mock_local_speed, "gateway_ping": mock_gateway_ping, "gateway_speed": mock_gateway_speed, - "cleanup": mock_cleanup, + "subprocess_run": mock_subproc_run, "log": mock_log, "chrome": mock_chrome, "access_code": mock_access_code, @@ -104,25 +104,8 @@ def test_perform_checks_gateway_speed_test_interval(mock_tasks, monkeypatch): perform_checks() # Run 1 mock_tasks["gateway_speed"].assert_not_called() - -def test_perform_checks_calls_cleanup_first(mock_tasks, monkeypatch): - """Ensure the pre-emptive cleanup runs at the beginning of perform_checks.""" - # Configure toggles so at least one downstream task runs - monkeypatch.setattr(config_module, "RUN_LOCAL_PING_TEST", True) - monkeypatch.setattr(config_module, "RUN_LOCAL_GATEWAY_PING_TEST", False) - monkeypatch.setattr(config_module, "RUN_LOCAL_SPEED_TEST", False) - - events: list[str] = [] - mock_tasks["cleanup"].side_effect = lambda *_, **__: events.append("cleanup") - mock_tasks["wifi"].side_effect = lambda: (events.append("wifi"), {})[1] - - perform_checks() - - # cleanup should be called and recorded before wifi diagnostics - assert mock_tasks["cleanup"].call_count == 1 - assert events and events[0] == "cleanup" - - # Do not assert interval behavior here; covered by dedicated test + # Note: Pre-emptive cleanup now occurs within the context manager. + # Dedicated tests cover that behavior in tests/test_cleanup.py. def test_perform_checks_gateway_speed_test_disabled(mock_tasks, monkeypatch):