From bb00f69fe58e92e1bcc651670fd31c80f4144ef7 Mon Sep 17 00:00:00 2001 From: OM222O Date: Mon, 22 Dec 2025 07:05:24 +0000 Subject: [PATCH] Improved typehints Simplified subprocess operations Cache result of re.compile instead of re-calculating it every time --- tests/test_utils.py | 33 ++++++++++++++------------------- wolnut/cli.py | 3 ++- wolnut/config.py | 15 +++++++-------- wolnut/monitor.py | 6 +++--- wolnut/utils.py | 29 +++++++++++++---------------- 5 files changed, 39 insertions(+), 47 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 1bbf016..dd0f83c 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,31 +1,26 @@ -import subprocess - from wolnut import utils +from unittest.mock import patch, Mock def test_validate_mac_format(): valid_mac_address = "de:ad:be:ef:be:ad" - invalid_mac_address = "ea:ts:be:ef:be:ad" # Mmm. Delicious beef bead. + invalid_mac_address = "ea:ts:be:ef:be:ad" - # `validate_mac_format()` returns a bool, so no need to compare with `==` assert utils.validate_mac_format(valid_mac_address) assert not utils.validate_mac_format(invalid_mac_address) def test_resolve_mac_from_host(mocker): # [ TODO - Issue #24 ] - Write tests that assert the appropriate exceptions were raised - class MockSubprocessResultOkay(object): - stdout = "de:ad:be:ef:be:ad" - - mock_subprocess_run = mocker.patch("wolnut.utils.subprocess.run") - mock_subprocess_run.return_value = MockSubprocessResultOkay() - result = utils.resolve_mac_from_host("localhost") - assert result == MockSubprocessResultOkay.stdout - mock_subprocess_run.assert_any_call( - ["ping", "-c", "1", "localhost"], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - mock_subprocess_run.assert_any_call( - ["arp", "-n", "localhost"], capture_output=True, text=True - ) + EXPECTED_RESULT = "de:ad:be:ef:be:ad" + EXPECTED_HOST = "localhost" + with patch("wolnut.utils.Popen", Mock()) as popen_mock: + popen_mock().communicate.return_value = EXPECTED_RESULT + popen_mock.call_args_list.clear() + assert utils.resolve_mac_from_host(EXPECTED_HOST) == EXPECTED_RESULT + assert len(popen_mock.call_args_list) == 2 + ping_call_args, arp_call_args = popen_mock.call_args_list + command, *_ , destination = ping_call_args.args[0] + assert (command, destination) == ("ping", EXPECTED_HOST) + command, *_ , destination = arp_call_args.args[0] + assert (command, destination) == ("arp", EXPECTED_HOST) \ No newline at end of file diff --git a/wolnut/cli.py b/wolnut/cli.py index 53a0196..b0ff736 100644 --- a/wolnut/cli.py +++ b/wolnut/cli.py @@ -7,6 +7,7 @@ from wolnut.state import ClientStateTracker from wolnut.monitor import get_ups_status, is_client_online from wolnut.wol import send_wol_packet +from typing import Optional logger = logging.getLogger("wolnut") @@ -196,7 +197,7 @@ def main(config_file: str, status_file: str, verbose: bool = False) -> int: help="The status filepath to load. Can also be set with WOLNUT_STATUS_FILE env var.", ) @click.option("--verbose", is_flag=True, help="Enable verbose logging") -def wolnut(config_file: str | None, status_file: str | None, verbose: bool) -> int: +def wolnut(config_file: Optional[str], status_file: Optional[str], verbose: bool) -> int: """A service to send Wake-on-LAN packets to clients after a power outage.""" logging.basicConfig( level=logging.INFO, diff --git a/wolnut/config.py b/wolnut/config.py index bb5f604..b8a455a 100644 --- a/wolnut/config.py +++ b/wolnut/config.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from pathlib import Path -from typing import Optional +from typing import Optional, Dict, Any from wolnut.state import DEFAULT_STATE_FILEPATH from wolnut.utils import validate_mac_format, resolve_mac_from_host @@ -19,8 +19,8 @@ class NutConfig: ups: str port: int = 3493 timeout: int = 5 - username: str | None = None - password: str | None = None + username: Optional[str] = None + password: Optional[str] = None @dataclass @@ -65,7 +65,7 @@ def find_state_file(state_file: Optional[str] = None) -> str: def load_config( - config_path: str, status_path: str = None, verbose: bool = False + config_path: str, status_path: Optional[str] = None, verbose: bool = False ) -> Optional[WolnutConfig]: try: with open(config_path, "r") as f: @@ -73,11 +73,10 @@ def load_config( validate_config(raw) except FileNotFoundError: logger.error("Config file not found at '%s'.", config_path) - return None + return except Exception: logger.exception("Failed to load or parse config file: '%s'.\n", config_path) - return None - + return # LOGGING... nut = NutConfig(**raw["nut"]) @@ -126,7 +125,7 @@ def load_config( return wolnut_config -def validate_config(raw: dict): +def validate_config(raw: Dict[str,Any]) -> None: if "clients" not in raw or not isinstance(raw["clients"], list): raise ValueError("Missing or invalid 'clients' list") diff --git a/wolnut/monitor.py b/wolnut/monitor.py index 9af8763..3ebdb37 100644 --- a/wolnut/monitor.py +++ b/wolnut/monitor.py @@ -1,14 +1,14 @@ import subprocess import logging import platform -from typing import Optional +from typing import Optional, Dict logger = logging.getLogger("wolnut") def get_ups_status( ups_name: str, username: Optional[str] = None, password: Optional[str] = None -) -> dict: +) -> Dict[str,str]: env = None if username and password: @@ -54,4 +54,4 @@ def is_client_online(host: str) -> bool: return result.returncode == 0 except Exception as e: logger.warning("Failed to ping %s: %s", host, e) - return False + return False diff --git a/wolnut/utils.py b/wolnut/utils.py index 6babace..7989ea9 100644 --- a/wolnut/utils.py +++ b/wolnut/utils.py @@ -1,7 +1,9 @@ -import subprocess +from subprocess import Popen, PIPE import re import logging +from typing import Optional +_MAC_ADDRESS_PATTERN = re.compile(r"([0-9A-Fa-f]{2}[:\-]){5}([0-9A-Fa-f]{2})") logger = logging.getLogger("wolnut") @@ -15,11 +17,10 @@ def validate_mac_format(mac: str) -> bool: Returns: bool: True if valid. """ - pattern = re.compile(r"^([0-9A-Fa-f]{2}[:\-]){5}([0-9A-Fa-f]{2})$") - return bool(pattern.match(mac)) + return _MAC_ADDRESS_PATTERN.fullmatch(mac) is not None -def resolve_mac_from_host(host: str) -> str | None: +def resolve_mac_from_host(host: str) -> Optional[str]: """ Attempts to resolve MAC address with provided hostname or IP. @@ -27,27 +28,23 @@ def resolve_mac_from_host(host: str) -> str | None: host (str): IP or Hostname. Returns: - str | None: The MAC address as a colon-separated string if found, + Optional[str]: The MAC address as a colon-separated string if found, otherwise None. """ + # FIXME: is it really necessary to ping before running ARP? + # if cache is not populated, it should be resolved correctly; if not this hints at a deeper bug + # Send a ping to ensure the ARP cache is populated try: - subprocess.run( - ["ping", "-c", "1", host], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) + Popen(["ping", "-c", "1", host]).communicate() except Exception as e: logger.warning("Failed to ping: %s: %s", host, e) - return None + return # Read the ARP table try: - result = subprocess.run(["arp", "-n", host], capture_output=True, text=True) - match = re.search(r"(([0-9A-Fa-f]{2}[:\-]){5}[0-9A-Fa-f]{2})", result.stdout) - if match: + out = Popen(["arp", "-n", host], text=True, stdout=PIPE).communicate() + if (match := _MAC_ADDRESS_PATTERN.search(out)): return match.group(0) except Exception as e: logger.warning("Failed to resolve ARP for: %s: %s", host, e) - - return None