diff --git a/custom_components/protocol_wizard/__init__.py b/custom_components/protocol_wizard/__init__.py index 1b43408..ab2f71d 100644 --- a/custom_components/protocol_wizard/__init__.py +++ b/custom_components/protocol_wizard/__init__.py @@ -416,7 +416,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) elif protocol_name == CONF_PROTOCOL_SNMP: - client = _create_snmp_client(config) + client = _create_snmp_client(config, hass) elif protocol_name == CONF_PROTOCOL_MQTT: client = _create_mqtt_client(config) else: @@ -600,7 +600,7 @@ async def _create_modbus_client(hass: HomeAssistant, config: dict, entry: Config slave_id, id(pymodbus_client)) return ModbusClient(pymodbus_client, slave_id) -def _create_snmp_client(config: dict) -> SNMPClient: +def _create_snmp_client(config: dict, hass: HomeAssistant | None = None) -> SNMPClient: """Create SNMP client (no caching needed - connectionless).""" from .protocols.snmp import SNMPClient @@ -609,6 +609,7 @@ def _create_snmp_client(config: dict) -> SNMPClient: port=config.get(CONF_PORT, 161), community=config.get("community", "public"), version=config.get("version", "2c"), + hass=hass, ) def _create_mqtt_client(config: dict) -> MQTTClient: diff --git a/custom_components/protocol_wizard/config_flow.py b/custom_components/protocol_wizard/config_flow.py index 61c9a9e..62a870f 100644 --- a/custom_components/protocol_wizard/config_flow.py +++ b/custom_components/protocol_wizard/config_flow.py @@ -818,6 +818,7 @@ async def _async_test_snmp_connection(self, data: dict[str, Any]) -> None: port=data.get(CONF_PORT, 161), community=data["community"], version=data["version"], + hass=self.hass, ) try: diff --git a/custom_components/protocol_wizard/protocols/snmp/client.py b/custom_components/protocol_wizard/protocols/snmp/client.py index 9e6c434..adbf38d 100644 --- a/custom_components/protocol_wizard/protocols/snmp/client.py +++ b/custom_components/protocol_wizard/protocols/snmp/client.py @@ -4,6 +4,7 @@ import asyncio import logging +import os from typing import Any from pysnmp.hlapi.v3arch.asyncio import ( @@ -22,6 +23,81 @@ _LOGGER = logging.getLogger(__name__) +# A single PySNMP engine is shared by every SNMP client in this process. Creating +# one is expensive (it loads MIB modules from disk), so it is built once inside an +# executor and kept for the lifetime of the process. +_SNMP_ENGINE: SnmpEngine | None = None +_ENGINE_LOCK = asyncio.Lock() + + +def _iter_pysnmp_mib_module_names(mib_builder): + """Yield bundled PySNMP MIB module names from file-backed MIB sources.""" + seen = set() + + for mib_source in mib_builder.get_mib_sources(): + source_dir = getattr(mib_source, "_srcName", None) + if not source_dir or not os.path.isdir(source_dir): + continue + + for _root, _dirs, files in os.walk(source_dir): + for filename in files: + if not filename.endswith(".py") or filename == "__init__.py": + continue + + module_name = os.path.splitext(filename)[0] + if module_name not in seen: + seen.add(module_name) + yield module_name + + +def _create_engine() -> SnmpEngine: + """Create and warm PySNMP's engine before it is used on the event loop.""" + engine = SnmpEngine() + + # PySNMP lazily loads bundled MIB modules during the first request and + # response. Home Assistant flags those file reads when they happen on the + # event loop, so force the lazy work into the executor with engine creation. + mib_builder = engine.get_mib_builder() + module_names = [ + "SNMPv2-SMI", + "SNMPv2-TC", + "SNMPv2-CONF", + "SNMPv2-TM", + "SNMPv2-MIB", + "PYSNMP-SOURCE-MIB", + "__SNMPv2-MIB", + ] + module_names.extend(_iter_pysnmp_mib_module_names(mib_builder)) + + for module_name in dict.fromkeys(module_names): + try: + mib_builder.load_modules(module_name) + except Exception as err: + _LOGGER.debug("Unable to preload PySNMP MIB module %s: %s", module_name, err) + + try: + mib_builder.import_symbols("SNMPv2-MIB", "snmpInPkts", "snmpOutPkts") + except Exception as err: + _LOGGER.debug("Unable to preload PySNMP MIB symbols: %s", err) + + return engine + + +async def _async_get_shared_engine(hass=None) -> SnmpEngine: + """Return the process-wide SNMP engine, creating it off the event loop.""" + global _SNMP_ENGINE + + async with _ENGINE_LOCK: + if _SNMP_ENGINE is None: + if hass is not None: + _SNMP_ENGINE = await hass.async_add_executor_job(_create_engine) + else: + loop = asyncio.get_running_loop() + _SNMP_ENGINE = await loop.run_in_executor(None, _create_engine) + _LOGGER.debug("SNMP engine created") + + return _SNMP_ENGINE + class SNMPClient(BaseProtocolClient): """SNMP client using pysnmp asyncio v3arch.""" @@ -34,7 +110,9 @@ def __init__( version: str = "2c", timeout: float = 5.0, retries: int = 3, + hass=None, ): + self.hass = hass self.host = host self.port = port self.community = community @@ -57,10 +135,12 @@ def __init__( self._context = ContextData() async def _ensure_engine(self) -> None: - """Lazily create engine and transport.""" + """Lazily attach the shared engine and create this client's transport.""" async with self._engine_lock: if self._engine is None: - self._engine = SnmpEngine() + # Engine creation reads MIB files from disk, so it is done once in + # an executor rather than on Home Assistant's event loop. + self._engine = await _async_get_shared_engine(self.hass) self._transport = await UdpTransportTarget.create( (self.host, self.port), timeout=self.timeout, @@ -81,16 +161,14 @@ async def connect(self) -> bool: return False async def disconnect(self) -> None: - """Clean up SNMP engine.""" - if self._engine: - try: - self._engine.close_dispatcher() - except Exception as err: - _LOGGER.debug("Error closing SNMP dispatcher: %s", err) - finally: - self._engine = None - self._transport = None - self._connected = False + """Release this client's SNMP resources. + + The engine is shared process-wide, so its dispatcher is deliberately left + open here — closing it would break every other SNMP client still in use. + """ + self._engine = None + self._transport = None + self._connected = False async def read(self, address: str, **kwargs) -> Any | None: """Read a single OID.""" diff --git a/tests/test_snmp_client_engine.py b/tests/test_snmp_client_engine.py new file mode 100644 index 0000000..a911a82 --- /dev/null +++ b/tests/test_snmp_client_engine.py @@ -0,0 +1,203 @@ +"""Tests for the shared, pre-warmed SNMP engine. + +PySNMP loads bundled MIB modules from disk the first time an engine encodes a +request. Doing that on Home Assistant's event loop trips the "blocking call +inside the event loop" warnings, so the engine is created once in an executor +with its MIB modules preloaded. +""" +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock + +import pytest + +from custom_components.protocol_wizard.protocols.snmp import client as snmp_client + + +@pytest.fixture(autouse=True) +def _reset_shared_engine(): + """Keep the process-wide engine cache out of other tests.""" + snmp_client._SNMP_ENGINE = None + yield + snmp_client._SNMP_ENGINE = None + + +class _FakeHass: + """Minimal hass double that records executor usage.""" + + def __init__(self): + self.executor_calls = [] + + async def async_add_executor_job(self, func, *args): + self.executor_calls.append(func) + return func(*args) + + +class TestCreateEngine: + """Test _create_engine warm-up.""" + + def test_preloads_bundled_mib_modules(self, monkeypatch): + mib_builder = MagicMock() + mib_builder.get_mib_sources.return_value = [] + engine = MagicMock() + engine.get_mib_builder.return_value = mib_builder + monkeypatch.setattr(snmp_client, "SnmpEngine", lambda: engine) + + assert snmp_client._create_engine() is engine + + loaded = [call.args[0] for call in mib_builder.load_modules.call_args_list] + # The modules PySNMP would otherwise read lazily on the event loop. + for module_name in ("SNMPv2-SMI", "SNMPv2-MIB", "SNMPv2-TM", "PYSNMP-SOURCE-MIB"): + assert module_name in loaded + # Each module is loaded at most once. + assert len(loaded) == len(set(loaded)) + mib_builder.import_symbols.assert_called_once_with( + "SNMPv2-MIB", "snmpInPkts", "snmpOutPkts" + ) + + def test_unloadable_module_does_not_raise(self, monkeypatch): + mib_builder = MagicMock() + mib_builder.get_mib_sources.return_value = [] + mib_builder.load_modules.side_effect = Exception("no such module") + mib_builder.import_symbols.side_effect = Exception("no such symbol") + engine = MagicMock() + engine.get_mib_builder.return_value = mib_builder + monkeypatch.setattr(snmp_client, "SnmpEngine", lambda: engine) + + assert snmp_client._create_engine() is engine + + +class TestIterMibModuleNames: + """Test _iter_pysnmp_mib_module_names.""" + + def test_discovers_module_names_from_file_sources(self, tmp_path): + mib_dir = tmp_path / "mibs" + mib_dir.mkdir() + (mib_dir / "SNMPv2-MIB.py").write_text("") + (mib_dir / "IF-MIB.py").write_text("") + (mib_dir / "__init__.py").write_text("") + (mib_dir / "notes.txt").write_text("") + + source = MagicMock() + source._srcName = str(mib_dir) + mib_builder = MagicMock() + mib_builder.get_mib_sources.return_value = [source] + + names = set(snmp_client._iter_pysnmp_mib_module_names(mib_builder)) + assert names == {"SNMPv2-MIB", "IF-MIB"} + + def test_skips_non_file_sources(self, tmp_path): + missing = MagicMock() + missing._srcName = str(tmp_path / "does-not-exist") + zipped = MagicMock() + zipped._srcName = None + mib_builder = MagicMock() + mib_builder.get_mib_sources.return_value = [missing, zipped] + + assert list(snmp_client._iter_pysnmp_mib_module_names(mib_builder)) == [] + + +class TestSharedEngine: + """Test _async_get_shared_engine.""" + + async def test_created_in_executor_and_cached(self, monkeypatch): + engine = object() + monkeypatch.setattr(snmp_client, "_create_engine", lambda: engine) + hass = _FakeHass() + + first = await snmp_client._async_get_shared_engine(hass) + second = await snmp_client._async_get_shared_engine(hass) + + assert first is engine + assert second is engine + # Built exactly once, and never on the event loop. + assert len(hass.executor_calls) == 1 + + async def test_falls_back_to_loop_executor_without_hass(self, monkeypatch): + engine = object() + creating_thread = [] + + def _create(): + import threading + + creating_thread.append(threading.current_thread()) + return engine + + monkeypatch.setattr(snmp_client, "_create_engine", _create) + + result = await snmp_client._async_get_shared_engine() + + assert result is engine + import threading + + assert creating_thread[0] is not threading.current_thread() + + async def test_concurrent_callers_share_one_engine(self, monkeypatch): + calls = [] + + def _create(): + calls.append(1) + return object() + + monkeypatch.setattr(snmp_client, "_create_engine", _create) + + engines = await asyncio.gather( + *(snmp_client._async_get_shared_engine() for _ in range(5)) + ) + + assert len(calls) == 1 + assert len({id(e) for e in engines}) == 1 + + +class TestClientUsesSharedEngine: + """Test SNMPClient engine handling.""" + + async def test_ensure_engine_uses_shared_engine(self, monkeypatch): + engine = object() + monkeypatch.setattr(snmp_client, "_create_engine", lambda: engine) + + async def _create_transport(*args, **kwargs): + return MagicMock() + + monkeypatch.setattr( + snmp_client.UdpTransportTarget, "create", _create_transport + ) + + hass = _FakeHass() + first = snmp_client.SNMPClient(host="10.0.0.1", hass=hass) + second = snmp_client.SNMPClient(host="10.0.0.2", hass=hass) + + await first._ensure_engine() + await second._ensure_engine() + + assert first._engine is engine + assert second._engine is engine + assert first._transport is not second._transport + assert len(hass.executor_calls) == 1 + + async def test_disconnect_leaves_shared_engine_open(self, monkeypatch): + engine = MagicMock() + monkeypatch.setattr(snmp_client, "_create_engine", lambda: engine) + + async def _create_transport(*args, **kwargs): + return MagicMock() + + monkeypatch.setattr( + snmp_client.UdpTransportTarget, "create", _create_transport + ) + + keeper = snmp_client.SNMPClient(host="10.0.0.1") + leaver = snmp_client.SNMPClient(host="10.0.0.2") + await keeper._ensure_engine() + await leaver._ensure_engine() + + await leaver.disconnect() + + # Unloading one config entry must not tear down the engine others use. + engine.close_dispatcher.assert_not_called() + assert leaver._engine is None + assert leaver._transport is None + assert leaver.is_connected is False + assert keeper._engine is engine + assert await snmp_client._async_get_shared_engine() is engine