diff --git a/bitelab/__init__.py b/bitelab/__init__.py index 51c0509..cf49659 100644 --- a/bitelab/__init__.py +++ b/bitelab/__init__.py @@ -4,8 +4,10 @@ from functools import lru_cache, wraps from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request from fastapi.security import OAuth2PasswordBearer from httpx import AsyncClient, Auth +from mock import patch, AsyncMock, Mock from pydantic import BaseModel -from starlette.status import HTTP_200_OK, HTTP_404_NOT_FOUND, HTTP_401_UNAUTHORIZED +from starlette.status import HTTP_200_OK, HTTP_404_NOT_FOUND, \ + HTTP_401_UNAUTHORIZED from . import config from . import data @@ -29,18 +31,65 @@ def new_parse_socket_addr(domain, addr): 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() + + if type == 'bool': + if res == 'true': + return True + elif res == 'false': + return False + + raise RuntimeError('unknown results for bool: %s' % repr(res)) + + raise RuntimeError('unknown type: %s' % repr(type)) + +async def snmpset(host, oid, value): + return await _snmpwrapper('set', host, oid, value=value) + +class Attribute: + defattrname = None + +class Power(Attribute): + defattrname = 'power' + +class SNMPPower(Power): + def __init__(self, host, port): + self.host = host + self.port = port + + # Future - add caching + invalidation on set + async def getvalue(self): + return await snmpget(self.host, + 'pethPsePortAdminEnable.1.%d' % self.port, 'bool') + + async def setvalue(self, v): + pass + class BoardClassInfo(BaseModel): clsname: str arch: str class BoardImpl: - def __init__(self, name, cls): + def __init__(self, name, brdclass, options): self.name = name - self.brdclass = cls + self.brdclass = brdclass + self.options = options + self.attrmap = {} + for i in options: + self.attrmap[i.defattrname] = i + + self.attrcache = {} + + async def update(self): + for i in self.attrmap: + self.attrcache[i] = await self.attrmap[i].getvalue() @property def attrs(self): - return {} + return dict(self.attrcache) class Board(BaseModel): name: str @@ -54,19 +103,23 @@ class BoardManager(object): board_class_info = { 'cora-z7s': { 'clsname': 'cora-z7s', - 'arch': 'arm64-aarch64', + 'arch': 'arm-armv7', }, } # Naming scheme: # - # - boards = { - 'cora-1': BoardImpl('cora-1', 'cora-z7s'), - } + board_gen = [ + dict(name='cora-1', brdclass='cora-z7s', options=[ + SNMPPower(host='poe', port=2), + ]), + ] def __init__(self, settings): self._settings = settings + self.boards = dict(**{ x.name: x for x in + (BoardImpl(**y) for y in self.board_gen)}) def classes(self): return self.board_class_info @@ -122,7 +175,8 @@ def get_boardmanager(settings: config.Settings = Depends(get_settings)): oauth2_scheme = OAuth2PasswordBearer(tokenUrl='/nonexistent') -async def lookup_user(token: str = Depends(oauth2_scheme), data: data.DataWrapper = Depends(get_data)): +async def lookup_user(token: str = Depends(oauth2_scheme), + data: data.DataWrapper = Depends(get_data)): try: return (await data.APIKey.objects.get(key=token)).user except orm.exceptions.NoMatch: @@ -147,7 +201,11 @@ async def foo(user: str = Depends(lookup_user), @router.get('/board_info',response_model=Dict[str, Board]) async def foo(user: str = Depends(lookup_user), brdmgr: BoardManager = Depends(get_boardmanager)): - return brdmgr.boards + brds = brdmgr.boards + for i in brds: + await brds[i].update() + + return brds @router.get('/') async def foo(board_prio: dict = Depends(board_priority), @@ -221,7 +279,8 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase): self.get_settings_override self.app.dependency_overrides[get_data] = self.get_data_override - self.client = AsyncClient(app=self.app, base_url='http://testserver') + self.client = AsyncClient(app=self.app, + base_url='http://testserver') def tearDown(self): self.app = None @@ -250,17 +309,48 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase): auth=BiteAuth('thisisanapikey')) self.assertEqual(res.status_code, HTTP_200_OK) self.assertEqual(res.json(), { 'cora-z7s': BoardClassInfo(**{ - 'arch': 'arm64-aarch64', 'clsname': 'cora-z7s', }) }) + 'arch': 'arm-armv7', 'clsname': 'cora-z7s', }) }) + + @patch('asyncio.create_subprocess_exec') + async def test_snmpwrapper(self, cse): + proc = Mock() + proc.communicate = AsyncMock() + proc.communicate.return_value = 'false\n' + cse.return_value = proc + + r = await snmpget('somehost', 'snmpoid', 'bool') + + self.assertEqual(r, False) + + cse.assert_called_with('snmpget', '-Oqv', 'somehost', 'snmpoid') - async def test_board_info(self): + proc.communicate.return_value = 'true\n' + r = await snmpget('somehost', 'snmpoid', 'bool') + + self.assertEqual(r, True) + + @patch('bitelab.snmpget') + async def test_board_info(self, sg): + # that when snmpget returns False + sg.return_value = False + + # that getting the board info res = await self.client.get('/board_info', auth=BiteAuth('thisisanapikey')) + + # calls snmpget w/ the correct args + sg.assert_called_with('poe', 'pethPsePortAdminEnable.1.2', + 'bool') + + # that it is successful self.assertEqual(res.status_code, HTTP_200_OK) + + # and returns the correct data info = { 'cora-1': { 'name': 'cora-1', 'brdclass': 'cora-z7s', - 'attrs': {}, + 'attrs': { 'power': False }, }, } self.assertEqual(res.json(), info) diff --git a/setup.py b/setup.py index bb24e13..82c1fc2 100644 --- a/setup.py +++ b/setup.py @@ -25,6 +25,7 @@ setup( 'aiokq @ git+https://www.funkthat.com/gitea/jmg/aiokq.git' 'orm', 'databases[sqlite]', + 'mock', ], extras_require = { 'dev': [ 'coverage' ],