Simplify SNMP driver

Author: NickCao

packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py

@@ -1,9 +1,11 @@
-import asyncio
 import socket
+from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from enum import Enum, IntEnum
-from typing import Any, Dict, Tuple
+from typing import Any, Dict
 
+from anyio import Event, fail_after
+from jumpstarter_driver_power.driver import PowerInterface, PowerReading
 from pysnmp.carrier.asyncio.dgram import udp
 from pysnmp.entity import config, engine
 from pysnmp.entity.rfc3413 import cmdgen
@@ -36,7 +38,7 @@ class SNMPError(Exception):
 
 
 @dataclass(kw_only=True)
-class SNMPServer(Driver):
+class SNMPServer(PowerInterface, Driver):
     """SNMP Power Control Driver"""
 
     host: str = field()
@@ -117,7 +119,7 @@ def _setup_snmp(self):
     def client(cls) -> str:
         return "jumpstarter_driver_snmp.client.SNMPServerClient"
 
-    def _create_snmp_callback(self, result: Dict[str, Any], response_received: asyncio.Event):
+    def _create_snmp_callback(self, result: Dict[str, Any], response_received: Event):
         def callback(snmpEngine, sendRequestHandle, errorIndication, errorStatus, errorIndex, varBinds, cbCtx):
             self.logger.debug(f"Callback {errorIndication} {errorStatus} {errorIndex} {varBinds}")
             if errorIndication:
@@ -138,29 +140,17 @@ def callback(snmpEngine, sendRequestHandle, errorIndication, errorStatus, errorI
 
         return callback
 
-    def _setup_event_loop(self) -> Tuple[asyncio.AbstractEventLoop, bool]:
-        try:
-            loop = asyncio.get_running_loop()
-            return loop, False
-        except RuntimeError:
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            return loop, True
-
-    async def _run_snmp_dispatcher(self, snmp_engine: engine.SnmpEngine, response_received: asyncio.Event):
+    async def _run_snmp_dispatcher(self, snmp_engine: engine.SnmpEngine, response_received: Event):
         snmp_engine.open_dispatcher()
         await response_received.wait()
         snmp_engine.close_dispatcher()
 
-    def _snmp_set(self, state: PowerState):
+    async def _snmp_set(self, state: PowerState):
         result = {"success": False, "error": None}
-        response_received = asyncio.Event()
-        loop = None
-        created_loop = False
+        response_received = Event()
 
         try:
             self.logger.info(f"Sending power {state.name} command to {self.host}")
-            loop, created_loop = self._setup_event_loop()
             snmp_engine = self._setup_snmp()
             callback = self._create_snmp_callback(result, response_received)
             cmdgen.SetCommandGenerator().send_varbinds(
@@ -172,10 +162,10 @@ def _snmp_set(self, state: PowerState):
                 callback,
             )
 
-            dispatcher_task = loop.create_task(self._run_snmp_dispatcher(snmp_engine, response_received))
             try:
-                loop.run_until_complete(asyncio.wait_for(dispatcher_task, self.timeout))
-            except asyncio.TimeoutError:
+                with fail_after(self.timeout):
+                    await self._run_snmp_dispatcher(snmp_engine, response_received)
+            except TimeoutError:
                 self.logger.warning(f"SNMP operation timed out after {self.timeout} seconds")
                 result["error"] = "SNMP operation timed out"
 
@@ -188,19 +178,20 @@ def _snmp_set(self, state: PowerState):
             error_msg = f"SNMP set failed: {str(e)}"
             self.logger.error(error_msg)
             raise SNMPError(error_msg) from e
-        finally:
-            if created_loop and loop:
-                loop.close()
 
     @export
-    def on(self):
+    async def on(self):
         """Turn power on"""
-        return self._snmp_set(PowerState.ON)
+        return await self._snmp_set(PowerState.ON)
 
     @export
-    def off(self):
+    async def off(self):
         """Turn power off"""
-        return self._snmp_set(PowerState.OFF)
+        return await self._snmp_set(PowerState.OFF)
+
+    @export
+    async def read(self) -> AsyncGenerator[PowerReading, None]:
+        raise NotImplementedError
 
     def close(self):
         """No cleanup needed since engines are created per operation"""

packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py

@@ -55,7 +55,8 @@ def setup_mock_snmp_engine():
         },
     ],
 )
-def test_snmp_auth_configurations(auth_config):
+@pytest.mark.anyio
+async def test_snmp_auth_configurations(auth_config):
     """Test different SNMP authentication configurations"""
     with (
         patch("pysnmp.entity.config.add_v3_user") as mock_add_user,
@@ -103,15 +104,13 @@ def test_snmp_auth_configurations(auth_config):
 
 @patch("pysnmp.entity.config.add_v3_user")
 @patch("pysnmp.entity.engine.SnmpEngine")
-def test_power_on_command(mock_engine, mock_add_user):
+@pytest.mark.anyio
+async def test_power_on_command(mock_engine, mock_add_user):
     """Test power on command execution"""
     mock_engine.return_value = setup_mock_snmp_engine()
 
     with (
         patch("pysnmp.entity.rfc3413.cmdgen.SetCommandGenerator.send_varbinds") as mock_send,
-        patch("asyncio.get_running_loop", side_effect=RuntimeError),
-        patch("asyncio.new_event_loop"),
-        patch("asyncio.set_event_loop"),
         patch("pysnmp.entity.config.add_target_parameters"),
         patch("pysnmp.entity.config.add_target_address"),
         patch("pysnmp.entity.config.add_transport"),
@@ -124,22 +123,20 @@ def side_effect(*args):
 
         mock_send.side_effect = side_effect
 
-        result = server.on()
+        result = await server.on()
         assert "Power ON command sent successfully" in result
         mock_send.assert_called_once()
 
 
 @patch("pysnmp.entity.config.add_v3_user")
 @patch("pysnmp.entity.engine.SnmpEngine")
-def test_power_off_command(mock_engine, mock_add_user):
+@pytest.mark.anyio
+async def test_power_off_command(mock_engine, mock_add_user):
     """Test power off command execution"""
     mock_engine.return_value = setup_mock_snmp_engine()
 
     with (
         patch("pysnmp.entity.rfc3413.cmdgen.SetCommandGenerator.send_varbinds") as mock_send,
-        patch("asyncio.get_running_loop", side_effect=RuntimeError),
-        patch("asyncio.new_event_loop"),
-        patch("asyncio.set_event_loop"),
         patch("pysnmp.entity.config.add_target_parameters"),
         patch("pysnmp.entity.config.add_target_address"),
         patch("pysnmp.entity.config.add_transport"),
@@ -152,6 +149,6 @@ def side_effect(*args):
 
         mock_send.side_effect = side_effect
 
-        result = server.off()
+        result = await server.off()
         assert "Power OFF command sent successfully" in result
         mock_send.assert_called_once()