add code for snmpset

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)