diff --git a/bitelab/snmp.py b/bitelab/snmp.py index d9ee553..7706fb3 100644 --- a/bitelab/snmp.py +++ b/bitelab/snmp.py @@ -37,24 +37,43 @@ import unittest __all__ = [ 'snmpget', 'snmpset', 'SNMPPower' ] -async def snmpget(host, oid, type): +def _tosnmp(typ, value): + if typ == 'bool': + if value: + outv = 'true' + else: + outv = 'false' + + return ('i', outv) + + raise RuntimeError('unknown type: %s' % repr(typ)) + +def _fromsnmp(typ, value): + if typ == 'bool': + if value == b'true': + return True + elif value == b'false': + return False + + raise RuntimeError('unknown results for bool: %s' % repr(value)) + + raise RuntimeError('unknown type: %s' % repr(typ)) + +async def snmpget(host, oid, typ): p = await asyncio.create_subprocess_exec('snmpget', '-Oqv', host, oid, stdout=subprocess.PIPE) res = (await p.communicate())[0].strip() - if type == 'bool': - if res == b'true': - return True - elif res == b'false': - return False + return _fromsnmp(typ, res) - raise RuntimeError('unknown results for bool: %s' % repr(res)) +async def snmpset(host, oid, typ, value): + p = await asyncio.create_subprocess_exec('snmpset', '-Oqv', host, oid, + *_tosnmp(typ, value), stdout=subprocess.PIPE) - raise RuntimeError('unknown type: %s' % repr(type)) + res = (await p.communicate())[0].strip() -async def snmpset(host, oid, value): - return await _snmpwrapper('set', host, oid, value=value) + return _fromsnmp(typ, res) class SNMPPower(Power): def __init__(self, host, port): @@ -66,14 +85,46 @@ class SNMPPower(Power): return await snmpget(self.host, 'pethPsePortAdminEnable.1.%d' % self.port, 'bool') - async def setvalue(self, v): + async def setvalue(self, typ, v): return await snmpset(self.host, - 'pethPsePortAdminEnable.1.%d' % self.port, v) + 'pethPsePortAdminEnable.1.%d' % self.port, typ, v) class TestSNMPWrapper(unittest.IsolatedAsyncioTestCase): @patch('asyncio.create_subprocess_exec') - async def test_snmpwrapper(self, cse): + async def test_snmpset(self, cse): + # that when snmpset returns false + wrap_subprocess_exec(cse, b'false\n') + + # when being set to false + r = await snmpset('somehost', 'snmpoid', 'bool', False) + + # that it returns false + self.assertEqual(r, False) + + # and is called with the correct parameters + cse.assert_called_with('snmpset', '-Oqv', 'somehost', + 'snmpoid', 'i', 'false', stdout=subprocess.PIPE) + + # that when snmpset returns true + wrap_subprocess_exec(cse, b'true\n') + + # when being set to true + r = await snmpset('somehost', 'snmpoid', 'bool', True) + + # that it returns true + self.assertEqual(r, True) + + # and is called with the correct parameters + cse.assert_called_with('snmpset', '-Oqv', 'somehost', + 'snmpoid', 'i', 'true', stdout=subprocess.PIPE) + + async def test_snmpset_wrongtype(self): + with self.assertRaises(RuntimeError): + await snmpset('somehost', 'snmpoid', 'boi', None) + + @patch('asyncio.create_subprocess_exec') + async def test_snmpget(self, cse): wrap_subprocess_exec(cse, b'false\n') r = await snmpget('somehost', 'snmpoid', 'bool') @@ -115,8 +166,8 @@ class TestSNMPPower(unittest.IsolatedAsyncioTestCase): 'bool') # that when setvalue is called - await sp.setvalue(True) + await sp.setvalue('bool', True) # calls snmpset w/ the correct args ss.assert_called_with('host', 'pethPsePortAdminEnable.1.5', - True) + 'bool', True)