From c25d72a821caa181635e18a9de489f810b37cc04 Mon Sep 17 00:00:00 2001 From: GoWithitRoger Date: Mon, 11 Aug 2025 10:36:08 -0500 Subject: [PATCH 1/2] =?UTF-8?q?feat(main):=20add=20pre-emptive=20chromedri?= =?UTF-8?q?ver=20cleanup;=20improve=20Wi=E2=80=91Fi=20Tx=20rate=20parsing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - add cleanup_old_processes() to proactively kill lingering chromedriver via pkill - call cleanup_old_processes() at start of perform_checks() to prevent orphaned drivers - enhance run_wifi_diagnostics_task(): detect Tx rate across multiple wdutil keys (Tx Rate, TxRate, Last Tx Rate, Max PHY Rate) - keep robust managed_webdriver_session post-run cleanup as secondary safeguard Tests - test_cleanup.py: cover pkill invocation and exception handling - test_main_flow.py: mock cleanup_old_processes; assert it runs first in perform_checks - test_core.py: parametrize Wi‑Fi Tx rate key variants; maintain existing success/failure cases - minor import/lint cleanups in tests Notes - aligns with acceptance criteria: pre-emptive cleanup logs, no orphaned chromedriver, primary flows intact Refs Ticket: Implement Pre-emptive and Post-run Process Cleanup --- main.py | 164 ++++++++++++++++++++++---------- tests/test_cleanup.py | 25 +++++ tests/test_core.py | 28 ++++++ tests/test_main_flow.py | 26 +++-- tests/test_webdriver_context.py | 95 ++++++++++++++++++ 5 files changed, 279 insertions(+), 59 deletions(-) create mode 100644 tests/test_cleanup.py create mode 100644 tests/test_webdriver_context.py diff --git a/main.py b/main.py index 8033bdc..d44c5f5 100644 --- a/main.py +++ b/main.py @@ -15,6 +15,7 @@ import re import subprocess import time +from contextlib import contextmanager from datetime import datetime from typing import ClassVar, Literal, Mapping, Optional, TypedDict @@ -55,6 +56,29 @@ 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`. @@ -78,6 +102,56 @@ def log_running_chromedriver_processes(debug_logger: DebugLogger) -> None: debug_logger.log(f"Error while checking chromedriver processes: {e}") +@contextmanager +def managed_webdriver_session(chrome_options: Options, debug_logger: DebugLogger): + """A self-cleaning context manager for the Selenium WebDriver. + + Sets up the driver and guarantees the chromedriver service process is + terminated on exit, aggressively if needed to avoid orphans. + """ + service: ChromeService = ChromeService() + driver: Optional[WebDriver] = None + print("Setting up WebDriver for gateway tests...") + try: + 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}") + yield driver + finally: + print("Shutting down WebDriver session...") + if driver: + debug_logger.log("WebDriver quit: START") + try: + driver.quit() # Graceful shutdown of browser and driver + 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. " + f"Error: {e}" + ) + # Verify process state after teardown + log_running_chromedriver_processes(debug_logger) + + # --- Typing Models --- class GatewayPingResults(TypedDict, total=False): """Structured results from gateway ping parsing.""" @@ -454,9 +528,6 @@ def run_local_ping_task(target: str) -> LocalPingResults: return {} -# In main.py - - def run_local_speed_test_task() -> Optional[SpeedResults]: """ Runs a local speed test with a retry mechanism, returning numerical @@ -574,9 +645,20 @@ def find_value(key: str, text: str) -> str: results["wifi_rssi"] = find_value("RSSI", output) results["wifi_noise"] = find_value("Noise", output) - results["wifi_tx_rate"] = find_value("TxRate", output) results["wifi_channel"] = find_value("Channel", output) + # Try a list of possible keys for transmit rate to make it more universal + tx_rate_keys = ["Tx Rate", "TxRate", "Last Tx Rate", "Max PHY Rate"] + for key in tx_rate_keys: + tx_rate = find_value(key, output) + if tx_rate != "N/A": + # Found it, so we can stop looking + results["wifi_tx_rate"] = tx_rate + break + else: + # If the loop finishes without finding any key, default to N/A + results["wifi_tx_rate"] = "N/A" + except Exception as e: print(f"Warning: Could not parse wdutil output. Error: {e}") @@ -628,6 +710,8 @@ 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})..." @@ -676,56 +760,34 @@ def perform_checks() -> None: chrome_options.add_argument("--no-sandbox") chrome_options.add_argument("--disable-dev-shm-usage") - driver: Optional[WebDriver] = None - # Use Selenium Manager to handle the driver - service: ChromeService = ChromeService() try: debug_log.log("Selenium setup: START") - print("Setting up WebDriver for gateway tests...") - 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_log.set_chromedriver_pid(service.process.pid) - debug_log.log(f"WebDriver service started with PID: {service.process.pid}") - except Exception as e: - debug_log.log(f"Unable to retrieve WebDriver PID: {e}") - debug_log.log("Selenium setup: END") - driver.get(config.GATEWAY_URL) - time.sleep(2) - debug_log.log("run_ping_test_task: START") - gateway_ping_results = run_ping_test_task(driver) - 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) - debug_log.log("run_speed_test_task: END") - if gateway_speed_results: - master_results.update(gateway_speed_results) + with managed_webdriver_session(chrome_options, debug_log) as driver: + debug_log.log("Selenium setup: END") + driver.get(config.GATEWAY_URL) + time.sleep(2) + debug_log.log("run_ping_test_task: START") + gateway_ping_results = run_ping_test_task(driver) + 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) + debug_log.log("run_speed_test_task: END") + if gateway_speed_results: + master_results.update(gateway_speed_results) except Exception as e: - print(f"An error occurred during the gateway automation process: {e}") - if 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.") - finally: - if driver: - debug_log.log("WebDriver quit: START") - print("Closing WebDriver...") - driver.quit() - debug_log.log("WebDriver quit: END") - if service: - debug_log.log("WebDriver service.stop(): START") - service.stop() - debug_log.log("WebDriver service.stop(): END") - # Verify process state immediately after stopping service - log_running_chromedriver_processes(debug_log) + 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 debug_log.log("perform_checks: END") log_results(master_results) diff --git a/tests/test_cleanup.py b/tests/test_cleanup.py new file mode 100644 index 0000000..7930a26 --- /dev/null +++ b/tests/test_cleanup.py @@ -0,0 +1,25 @@ +import os +import sys +from unittest.mock import 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_cleanup_old_processes_handles_exception(): + with patch("main.subprocess.run", side_effect=RuntimeError("boom")): + # Should not raise + main.cleanup_old_processes() diff --git a/tests/test_core.py b/tests/test_core.py index 9f53542..5dbfb48 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -5,6 +5,8 @@ import sys from unittest.mock import MagicMock, patch +import pytest + # Ensure the main module can be imported sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) @@ -170,3 +172,29 @@ def test_wifi_diagnostics_failure(mock_run): assert results["wifi_rssi"] == "N/A" assert results["wifi_noise"] == "N/A" assert results["wifi_bssid"] == "N/A" + + +@patch("main.subprocess.run") +@pytest.mark.parametrize( + "tx_key", + [ + "Tx Rate", + "TxRate", + "Last Tx Rate", + "Max PHY Rate", + ], +) +def test_wifi_diagnostics_tx_rate_variants(mock_run, tx_key): + """Ensure tx rate is parsed correctly across multiple wdutil key variants.""" + # Build wdutil output with the variant key + wd_output = f"RSSI: -55\nNoise: -90\n{tx_key}: 866\nChannel: 149,80" + route_output = "gateway: 192.168.1.1" + arp_output = "? (192.168.1.1) at a1:b2:c3:d4:e5:f6 on en0 ifscope [ethernet]" + mock_run.side_effect = [ + MagicMock(stdout=wd_output, returncode=0, stderr=""), # wdutil + MagicMock(stdout=route_output, returncode=0, stderr=""), # route + MagicMock(stdout="", returncode=0, stderr=""), # ping + MagicMock(stdout=arp_output, returncode=0, stderr=""), # arp + ] + results = run_wifi_diagnostics_task() + assert results["wifi_tx_rate"] == "866" diff --git a/tests/test_main_flow.py b/tests/test_main_flow.py index b437cb1..0c82e8a 100644 --- a/tests/test_main_flow.py +++ b/tests/test_main_flow.py @@ -23,6 +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.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, @@ -37,6 +38,7 @@ def mock_tasks(): "local_speed": mock_local_speed, "gateway_ping": mock_gateway_ping, "gateway_speed": mock_gateway_speed, + "cleanup": mock_cleanup, "log": mock_log, "chrome": mock_chrome, "access_code": mock_access_code, @@ -102,17 +104,25 @@ def test_perform_checks_gateway_speed_test_interval(mock_tasks, monkeypatch): perform_checks() # Run 1 mock_tasks["gateway_speed"].assert_not_called() - mock_tasks["gateway_speed"].reset_mock() - # Should run on second call (run #2) - perform_checks() # Run 2 - mock_tasks["gateway_speed"].assert_called_once() +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) - mock_tasks["gateway_speed"].reset_mock() + events: list[str] = [] + mock_tasks["cleanup"].side_effect = lambda *_, **__: events.append("cleanup") + mock_tasks["wifi"].side_effect = lambda: (events.append("wifi"), {})[1] - # Should NOT run on third call (run #3) - perform_checks() # Run 3 - mock_tasks["gateway_speed"].assert_not_called() + 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 def test_perform_checks_gateway_speed_test_disabled(mock_tasks, monkeypatch): diff --git a/tests/test_webdriver_context.py b/tests/test_webdriver_context.py new file mode 100644 index 0000000..052b1b5 --- /dev/null +++ b/tests/test_webdriver_context.py @@ -0,0 +1,95 @@ +import time +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import MagicMock, mock_open, patch + +import pytest + +import config +import main + + +def _make_service_with_process(pid: int = 12345) -> MagicMock: + process = MagicMock() + process.pid = pid + service = MagicMock() + service.process = process + return service + + +def _make_driver() -> MagicMock: + driver = MagicMock() + return driver + + +def test_managed_webdriver_session_teardown_calls_quit_and_kill() -> None: + debug_logger = main.DebugLogger(start_time=time.time()) + + fake_service = _make_service_with_process() + fake_driver = _make_driver() + + with patch.object(main, "ChromeService", return_value=fake_service): + with patch.object(main.webdriver, "Chrome", return_value=fake_driver): + with patch.object(main, "log_running_chromedriver_processes") as mock_log_procs: + chrome_options = main.Options() + # Exercise context manager + with main.managed_webdriver_session(chrome_options, debug_logger) as d: + assert d is fake_driver + # After exit: ensure proper teardown + fake_driver.quit.assert_called_once() + fake_service.process.kill.assert_called_once() + fake_service.process.wait.assert_called_once() + mock_log_procs.assert_called_once_with(debug_logger) + + +def test_managed_webdriver_session_handles_missing_pid_gracefully() -> None: + debug_logger = main.DebugLogger(start_time=time.time()) + + fake_service = _make_service_with_process(pid=None) # type: ignore[arg-type] + fake_service.process.pid = None + fake_driver = _make_driver() + + with patch.object(main, "ChromeService", return_value=fake_service): + with patch.object(main.webdriver, "Chrome", return_value=fake_driver): + with patch.object(main, "log_running_chromedriver_processes"): + chrome_options = main.Options() + with main.managed_webdriver_session(chrome_options, debug_logger): + pass + # Still should attempt kill/wait even without a PID + fake_driver.quit.assert_called_once() + fake_service.process.kill.assert_called_once() + fake_service.process.wait.assert_called_once() + + +def test_perform_checks_uses_context_manager_once(monkeypatch: pytest.MonkeyPatch) -> None: + # Reduce side-effects by disabling optional tests in config + monkeypatch.setattr(config, "RUN_GATEWAY_SPEED_TEST_INTERVAL", 0, raising=False) + monkeypatch.setattr(config, "RUN_LOCAL_PING_TEST", False, raising=False) + monkeypatch.setattr(config, "RUN_LOCAL_GATEWAY_PING_TEST", False, raising=False) + monkeypatch.setattr(config, "RUN_LOCAL_SPEED_TEST", False, raising=False) + + # Stub out local tasks to avoid subprocess work + monkeypatch.setattr(main, "run_local_ping_task", lambda target: {}) + monkeypatch.setattr(main, "run_local_speed_test_task", lambda: {}) + monkeypatch.setattr(main, "run_wifi_diagnostics_task", lambda: {}) + + # Stub Selenium tasks called inside the context + monkeypatch.setattr(main, "run_ping_test_task", lambda driver: {}) + + calls = SimpleNamespace(count=0) + + @contextmanager + def fake_manager(chrome_options, debug_logger): # type: ignore[no-untyped-def] + calls.count += 1 + yield MagicMock() + + monkeypatch.setattr(main, "managed_webdriver_session", fake_manager) + + # Prevent file writes in log_results + monkeypatch.setattr(main.os.path, "exists", lambda p: False) + monkeypatch.setattr(main.os.path, "getsize", lambda p: 0) + + # Exercise perform_checks and ensure context manager used once, while mocking open + with patch("builtins.open", new_callable=mock_open): + main.perform_checks() + assert calls.count == 1 From 7e82870e0f7a3489286e6e3514bb82527fa6c844 Mon Sep 17 00:00:00 2001 From: GoWithitRoger Date: Mon, 11 Aug 2025 10:44:52 -0500 Subject: [PATCH 2/2] fix: strip Mbps unit from wifi speed values for consistent parsing --- main.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/main.py b/main.py index d44c5f5..4807ef6 100644 --- a/main.py +++ b/main.py @@ -641,7 +641,10 @@ def run_wifi_diagnostics_task() -> WifiDiagnostics: def find_value(key: str, text: str) -> str: """Helper to find values in the wdutil output using regex.""" match = re.search(rf"^\s*{key}\s*:\s*(.*)$", text, re.MULTILINE) - return match.group(1).strip() if match else "N/A" + if match: + # Strip trailing unit to normalize values like '864.0 Mbps' -> '864.0' + return match.group(1).strip().replace(" Mbps", "") + return "N/A" results["wifi_rssi"] = find_value("RSSI", output) results["wifi_noise"] = find_value("Noise", output)