Skip to content

refactor: move chromedriver cleanup into managed_webdriver_session co… #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
GoWithitRoger merged 1 commit into from
Aug 11, 2025
Merged

refactor: move chromedriver cleanup into managed_webdriver_session co… #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
135 changes: 54 additions & 81 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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`.

Expand All @@ -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)

Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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) ---
Expand Down Expand Up @@ -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)
Expand Down
49 changes: 35 additions & 14 deletions tests/test_cleanup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,25 +1,46 @@
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__), "..")))

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)
25 changes: 4 additions & 21 deletions tests/test_main_flow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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):
Expand Down