add tests for running script at reserve/release

bitelab/__init__.py

@@ -11,11 +11,13 @@ from starlette.responses import JSONResponse
 from starlette.status import HTTP_200_OK
 from starlette.status import HTTP_400_BAD_REQUEST, HTTP_401_UNAUTHORIZED, \
     HTTP_403_FORBIDDEN, HTTP_404_NOT_FOUND, HTTP_409_CONFLICT
+from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
 
 from . import config
 from .data import *
 
 import asyncio
+import json
 import orm
 import os
 import socket
@@ -39,7 +41,7 @@ tcp_server.parse_socket_addr = new_parse_socket_addr
 async def snmpget(host, oid, type):
 	p = await asyncio.create_subprocess_exec('snmpget', '-Oqv', host, oid)
 
-	res = (await p.communicate()).strip()
+	res = (await p.communicate())[0].strip()
 
 	if type == 'bool':
 		if res == 'true':
@@ -98,6 +100,9 @@ class BoardImpl:
 
 		self.attrcache = {}
 
+	def __repr__(self): #pragma: no cover
+		return repr(Board.from_orm(self))
+
 	async def reserve(self):
 		assert self.lock.locked() and not self.reserved
 
@@ -112,6 +117,14 @@ class BoardImpl:
 		for i in self.attrmap:
 			self.attrcache[i] = await self.attrmap[i].getvalue()
 
+	def add_info(self, d):
+		self.attrcache.update(d)
+
+	def clean_info(self):
+		# clean up attributes
+		for i in set(self.attrcache) - set(self.attrmap):
+			del self.attrcache[i]
+
 	@property
 	def attrs(self):
 		return dict(self.attrcache)
@@ -198,7 +211,7 @@ def get_settings(): # pragma: no cover
 # how to get coverage for this?
 @unhashable_lru()
 def get_data(settings: config.Settings = Depends(get_settings)): # pragma: no cover
-	print(repr(settings))
+	#print(repr(settings))
 	database = data.databases.Database('sqlite:///' + settings.db_file)
 	d = make_orm(self.database)
 	return d
@@ -247,21 +260,54 @@ async def get_board_info(board_id, user: str = Depends(lookup_user),
 async def reserve_board(board_id_or_class, user: str = Depends(lookup_user),
     brdmgr: BoardManager = Depends(get_boardmanager),
     brdlck: asyncio.Lock = Depends(get_board_lock),
+    settings: config.Settings = Depends(get_settings),
     data: data.DataWrapper = Depends(get_data)):
 	board_id = board_id_or_class
 	brd = brdmgr.boards[board_id]
 
 	async with brd.lock:
 		try:
-			await data.BoardStatus.objects.create(board=board_id, user=user)
+			obrdreq = await data.BoardStatus.objects.create(board=board_id,
+			    user=user)
+			# XXX - There is a bug in orm where the returned
+			# object has an incorrect board value
+			# see: https://github.com/encode/orm/issues/47
+			#assert obrdreq.board == board_id and \
+			#    obrdreq.user == user
+			brdreq = await data.BoardStatus.objects.get(board=board_id,
+			    user=user)
 			await brd.reserve()
 		# XXX - orm isn't doing it's job here
 		except sqlite3.IntegrityError:
 			raise BITEError(
 			    status_code=HTTP_409_CONFLICT,
-			    errobj=Error(error='Board currently reserved.', board=Board.from_orm(brd)),
+			    errobj=Error(error='Board currently reserved.',
+			    board=Board.from_orm(brd)),
 			)
 
+		# Initialize board
+		try:
+			sub = await asyncio.create_subprocess_exec(
+			    settings.setup_script, 'reserve', brd.name, user)
+			stdout, stderr = await sub.communicate()
+			if sub.returncode:
+				raise RuntimeError(sub.returncode, stderr)
+
+			brd.add_info(json.loads(stdout))
+		except Exception as e:
+			await brdreq.delete()
+			await brd.release()
+			if isinstance(e, RuntimeError):
+				retcode, stderr = e.args
+				raise BITEError(
+				    status_code=HTTP_500_INTERNAL_SERVER_ERROR,
+				    errobj=Error(error=
+					'Failed to init board, ret: %d, stderr: %s' %
+					(retcode, repr(stderr)),
+				    board=Board.from_orm(brd)),
+				)
+			raise
+
 	await brd.update()
 
 	return brd
@@ -270,6 +316,7 @@ async def reserve_board(board_id_or_class, user: str = Depends(lookup_user),
 async def release_board(board_id, user: str = Depends(lookup_user),
     brdmgr: BoardManager = Depends(get_boardmanager),
     brdlck: asyncio.Lock = Depends(get_board_lock),
+    settings: config.Settings = Depends(get_settings),
     data: data.DataWrapper = Depends(get_data)):
 	brd = brdmgr.boards[board_id]
 
@@ -279,18 +326,26 @@ async def release_board(board_id, user: str = Depends(lookup_user),
 			if user != brduser.user:
 				raise BITEError(
 				    status_code=HTTP_403_FORBIDDEN,
-				    errobj=Error(error='Board reserved by %s.' % repr(brduser),
+				    errobj=Error(error='Board reserved by %s.' % repr(brduser.user),
 				    board=Board.from_orm(brd)))
 
-			await data.BoardStatus.delete(brduser)
-			await brd.release()
-
 		except orm.exceptions.NoMatch:
 			raise BITEError(
 			    status_code=HTTP_400_BAD_REQUEST,
 			    errobj=Error(error='Board not reserved.', board=Board.from_orm(brd)),
 			)
 
+		sub = await asyncio.create_subprocess_exec(
+		    settings.setup_script, 'release', brd.name, user)
+		stdout, stderr = await sub.communicate()
+		if sub.returncode:
+			raise RuntimeError(sub.returncode, stderr)
+
+		await data.BoardStatus.delete(brduser)
+		await brd.release()
+
+		brd.clean_info()
+
 	await brd.update()
 
 	return brd
@@ -378,6 +433,11 @@ async def _setup_data(data):
 	await data.APIKey.objects.create(user='foo', key='thisisanapikey')
 	await data.APIKey.objects.create(user='bar', key='anotherlongapikey')
 
+# Per RFC 5737 (https://tools.ietf.org/html/rfc5737):
+#   The blocks 192.0.2.0/24 (TEST-NET-1), 198.51.100.0/24 (TEST-NET-2),
+#   and 203.0.113.0/24 (TEST-NET-3) are provided for use in
+#   documentation.
+
 # Note: this will not work under python before 3.8 before
 # IsolatedAsyncioTestCase was added.  The tearDown has to happen
 # with the event loop running, otherwise the task and other things
@@ -446,12 +506,19 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase):
 		self.assertEqual(res.json(),  { 'cora-z7s': BoardClassInfo(**{
 		    'arch': 'arm-armv7', 'clsname': 'cora-z7s', }) })
 
-	@patch('asyncio.create_subprocess_exec')
-	async def test_snmpwrapper(self, cse):
+	@staticmethod
+	def _wrap_subprocess_exec(mockobj, stdout='', stderr='', retcode=0):
 		proc = Mock()
 		proc.communicate = AsyncMock()
-		proc.communicate.return_value = 'false\n'
-		cse.return_value = proc
+		proc.communicate.return_value = (stdout, stderr)
+		proc.wait = AsyncMock()
+		proc.wait.return_value = retcode
+		proc.returncode = retcode
+		mockobj.return_value = proc
+
+	@patch('asyncio.create_subprocess_exec')
+	async def test_snmpwrapper(self, cse):
+		self._wrap_subprocess_exec(cse, 'false\n')
 
 		r = await snmpget('somehost', 'snmpoid', 'bool')
 
@@ -459,13 +526,13 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase):
 
 		cse.assert_called_with('snmpget', '-Oqv', 'somehost', 'snmpoid')
 
-		proc.communicate.return_value = 'true\n'
+		self._wrap_subprocess_exec(cse, 'true\n')
 		r = await snmpget('somehost', 'snmpoid', 'bool')
 
 		self.assertEqual(r, True)
 
 		# that a bogus return value
-		proc.communicate.return_value = 'bogus\n'
+		self._wrap_subprocess_exec(cse, 'bogus\n')
 
 		# raises an error
 		with self.assertRaises(RuntimeError):
@@ -475,8 +542,9 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase):
 		with self.assertRaises(RuntimeError):
 			await snmpget('somehost', 'snmpoid', 'randomtype')
 
+	@patch('asyncio.create_subprocess_exec')
 	@patch('bitelab.snmpget')
-	async def test_board_reserve_release(self, sg):
+	async def test_board_reserve_release(self, sg, cse):
 		# that when releasing a board that is not yet reserved
 		res = await self.client.post('/board/cora-1/release',
 		    auth=BiteAuth('anotherlongapikey'))
@@ -487,6 +555,31 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase):
 		# that when snmpget returns False
 		sg.return_value = False
 
+		# that when the setup script will fail
+		self._wrap_subprocess_exec(cse, stderr='error', retcode=1)
+
+		# that reserving the board
+		res = await self.client.post('/board/cora-1/reserve',
+		    auth=BiteAuth('thisisanapikey'))
+
+		# that it is a failure
+		self.assertEqual(res.status_code, HTTP_500_INTERNAL_SERVER_ERROR)
+
+		# and returns the correct data
+		info = Error(error='Failed to init board, ret: 1, stderr: \'error\'',
+			board=Board(name='cora-1',
+				brdclass='cora-z7s',
+				reserved=False,
+			),
+		).dict()
+		self.assertEqual(res.json(), info)
+
+		# and that it called the start script
+		cse.assert_called_with(self.settings.setup_script, 'reserve', 'cora-1', 'foo')
+
+		# that when the setup script returns
+		self._wrap_subprocess_exec(cse, json.dumps(dict(ip='192.0.2.10')))
+
 		# that reserving the board
 		res = await self.client.post('/board/cora-1/reserve',
 		    auth=BiteAuth('thisisanapikey'))
@@ -495,26 +588,29 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase):
 		self.assertEqual(res.status_code, HTTP_200_OK)
 
 		# and returns the correct data
-		info = {
-			'name': 'cora-1',
-			'brdclass': 'cora-z7s',
-			'reserved': True,
-			'attrs': { 'power': False },
-		}
-		self.assertEqual(res.json(), info)
+		brdinfo = Board(name='cora-1',
+			brdclass='cora-z7s',
+			reserved=True,
+			attrs=dict(power=False,
+				ip='192.0.2.10',
+			),
+		).dict()
+		self.assertEqual(res.json(), brdinfo)
+
+		# and that it called the start script
+		cse.assert_called_with(self.settings.setup_script, 'reserve', 'cora-1', 'foo')
 
 		# that another user reserving the board
 		res = await self.client.post('/board/cora-1/reserve',
 		    auth=BiteAuth('anotherlongapikey'))
 
-		# that the request is successful
-		# it should likely be this, but can't get FastAPI exceptions to work
+		# that the request is fails with a conflict
 		self.assertEqual(res.status_code, HTTP_409_CONFLICT)
 
 		# and returns the correct data
 		info = {
 			'error': 'Board currently reserved.',
-			'board': info,
+			'board': brdinfo,
 		}
 		self.assertEqual(res.json(), info)
 
@@ -527,9 +623,10 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase):
 
 		# and returns the correct data
 		info = {
-			'error': 'Board not reserved by you.',
-			'board': info,
+			'error': 'Board reserved by \'foo\'.',
+			'board': brdinfo,
 		}
+		self.assertEqual(res.json(), info)
 
 		# that when the correct user releases the board
 		res = await self.client.post('/board/cora-1/release',
@@ -545,6 +642,10 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase):
 			'reserved': False,
 			'attrs': { 'power': False },
 		}
+		self.assertEqual(res.json(), info)
+
+		# and that it called the release script
+		cse.assert_called_with(self.settings.setup_script, 'release', 'cora-1', 'foo')
 
 		# that it can be reserved by a different user
 		res = await self.client.post('/board/cora-1/reserve',

bitelab/data.py

@@ -1,6 +1,6 @@
 from typing import Optional, Union, Dict, Any
 import databases
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from datetime import datetime
 import orm
 import sqlalchemy
@@ -15,7 +15,7 @@ class Board(BaseModel):
 	name: str
 	brdclass: str
 	reserved: bool
-	attrs: Dict[str, Any]
+	attrs: Dict[str, Any] = Field(default_factory=dict)
 
 	class Config:
 		orm_mode = True