class BackendFixture: """A test fixture for running a server and imperatively displaying views This fixture is typically used alongside async web drivers like ``playwight``. Example: .. code-block:: async with BackendFixture() as server: server.mount(MyComponent) """ _records: list[logging.LogRecord] _server_future: asyncio.Task[Any] _exit_stack = AsyncExitStack() def __init__( self, host: str = "127.0.0.1", port: Optional[int] = None, app: Any | None = None, implementation: BackendImplementation[Any] | None = None, options: Any | None = None, ) -> None: self.host = host self.port = port or find_available_port( host, allow_reuse_waiting_ports=False) self.mount, self._root_component = hotswap() if app is not None: if implementation is None: raise ValueError( "If an application instance its corresponding " "server implementation must be provided too.") self._app = app self.implementation = implementation or default_server self._options = options @property def log_records(self) -> list[logging.LogRecord]: """A list of captured log records""" return self._records def url(self, path: str = "", query: Optional[Any] = None) -> str: """Return a URL string pointing to the host and point of the server Args: path: the path to a resource on the server query: a dictionary or list of query parameters """ return urlunparse([ "http", f"{self.host}:{self.port}", path, "", urlencode(query or ()), "", ]) def list_logged_exceptions( self, pattern: str = "", types: Union[Type[Any], Tuple[Type[Any], ...]] = Exception, log_level: int = logging.ERROR, del_log_records: bool = True, ) -> list[BaseException]: """Return a list of logged exception matching the given criteria Args: log_level: The level of log to check exclude_exc_types: Any exception types to ignore del_log_records: Whether to delete the log records for yielded exceptions """ return list_logged_exceptions( self.log_records, pattern, types, log_level, del_log_records, ) async def __aenter__(self) -> BackendFixture: self._exit_stack = AsyncExitStack() self._records = self._exit_stack.enter_context(capture_idom_logs()) app = self._app or self.implementation.create_development_app() self.implementation.configure(app, self._root_component, self._options) started = asyncio.Event() server_future = asyncio.create_task( self.implementation.serve_development_app(app, self.host, self.port, started)) async def stop_server() -> None: server_future.cancel() try: await asyncio.wait_for(server_future, timeout=3) except asyncio.CancelledError: pass self._exit_stack.push_async_callback(stop_server) try: await asyncio.wait_for(started.wait(), timeout=3) except Exception: # pragma: no cover # see if we can await the future for a more helpful error await asyncio.wait_for(server_future, timeout=3) raise return self async def __aexit__( self, exc_type: Optional[Type[BaseException]], exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> None: await self._exit_stack.aclose() self.mount(None) # reset the view logged_errors = self.list_logged_exceptions(del_log_records=False) if logged_errors: # pragma: no cover raise LogAssertionError( "Unexpected logged exception") from logged_errors[0] return None
class EdgeOS: username = None password = None url = None ssl = None session = None tasks = {} stack = None def __init__(self, username, password, url, ssl=None, session_id=None): self.username = username self.password = password self.url = url self.ssl = ssl self.session = None self.session_id = None self.headers = {'Content-type': 'application/json'} self.cookies = {} self.sysdata = {'ping-data': {}} if session_id: self.session_id = session_id self.cookies['beaker.session.id'] = session_id async def setup(self): self.stack = AsyncExitStack() self.stack.push_async_callback(EdgeOS.close, self) if not self.session: # If no session provided create our own, push on stack to make sure it's cleaned up later self.session = await self.stack.enter_async_context( ClientSession(cookie_jar=CookieJar(unsafe=True), raise_for_status=True, timeout=ClientTimeout(15))) ''' Start login loop, won't return until first run is complete ''' await self.stack.enter_async_context( TaskEvery(stay_logged_in, self, interval=300)) async def close(self): for key, value in self.tasks.items(): try: value.cancel() await value except Exception as e: logger.warning(f"error canceling {key} got exception {e}") async def __aenter__(self): await self.setup() return self async def __aexit__(self, exception_type, exception_value, traceback): await self.stack.aclose() async def add_task(self, name, task): if self.tasks.get(name, None): logger.debug(f"PREMATURELY CLOSING {name}") old = self.tasks[name] try: old.cancel() await old except Exception as e: logger.warning(f"WHILE CANCELING {name} got exception {e}") self.tasks[name] = task async def is_logged_in(self): with suppress(): if await self.sys_info() != None: return True return False async def login(self): if self.username and self.password: async with self.session.post(f'{self.url}', data={ 'username': self.username, 'password': self.password }, ssl=self.ssl) as resp: try: s_id = self.session.cookie_jar.filter_cookies( self.url)['beaker.session.id'].value token = self.session.cookie_jar.filter_cookies( self.url)['X-CSRF-TOKEN'].value self.headers = {'X-CSRF-TOKEN': token} self.session_id = s_id self.cookies['beaker.session.id'] = s_id logger.debug("New seesion ending in {}".format(s_id[:4])) except Exception as e: logging.error("Failed to login") raise Exception(f"LOGIN ERROR {e!r}") async def data(self, data_type): try: result = None async with self.session.get( f"{self.url}/api/edge/data.json?data={data_type}", ssl=self.ssl, headers=self.headers, raise_for_status=False, cookies=self.cookies) as resp: if resp.status != 200: return None result = await resp.json() if result.get('success', 0): self.sysdata[data_type] = result.get('output', None) else: return None return result except Exception as e: logging.debug(f"exception in data {e!r}") return None async def data_every(self, data_type, interval): # self- registering task task = asyncio.current_task() if task: await self.add_task(data_type, task) with suppress(asyncio.CancelledError): while True: await self.data(data_type) await asyncio.sleep(interval) async def sys_info(self): return await self.data('sys_info') async def dhcp_leases(self): return await self.data('dhcp_leases') async def dhcp_stats(self): return await self.data('dhcp_stats') async def routes(self): return await self.data('routes') async def sys_info(self): return await self.data('sys_info') async def config(self): async with self.session.get(f"{self.url}/api/edge/get.json", ssl=self.ssl, headers=self.headers, cookies=self.cookies) as resp: if resp.status != 200: return None temp = await resp.json() if temp.get('success', False): self.sysconfig = temp['GET'] return temp['GET'] return None async def ping(self, target='1.1.1.1', count=3, size=100): ret = '' init = { 'SUBSCRIBE': [{ 'name': 'ping-feed', 'sub_id': f'ping-{target}', 'target': target, 'count': count, 'size': size }] } async for payload in self._ws(init=init, keepalive=False, timeout=15): ret += payload[f'ping-{target}'] if f"--- {target} ping statistics ---" in ret: self.sysdata['pinglast'] = ret return self.process_ping(ret, target) def process_ping(self, output, target): data = {target: {'time': time()}} for line in output.splitlines(): if 'packets transmitted' in line: sent, _, _, recv, *_ = line.split() data[target]['sent'] = int(sent) data[target]['lost'] = int(sent) - int(recv) continue if 'min/avg/max/mdev' in line: _, dat = line.split("=") dat = dat.split()[0] pdat = [float(x) for x in dat.split('/')] data[target]['min'] = pdat[0] data[target]['avg'] = pdat[1] data[target]['max'] = pdat[2] data[target]['mdev'] = pdat[3] self.sysdata['ping-data'].update(data) return data async def ping_every(self, interval=120, target='1.1.1.1', **kwargs): # self- registering task task = asyncio.current_task() if task: await self.add_task(f'ping-{target}-every', task) with suppress(asyncio.CancelledError): while True: await self.ping(target, **kwargs) await asyncio.sleep(interval) async def background_stats(self, subs=[ "export", "discover", "interfaces", "system-stats", "num-routes", "config-change", "users" ]): # self- registering task task = asyncio.current_task() if task: await self.add_task('stats', task) with suppress(asyncio.CancelledError): async for payload in self.stats(subs): pass async def stats(self, subs=[ "export", "discover", "interfaces", "system-stats", "num-routes", "config-change", "users" ], reload_on_change=True): if reload_on_change and (not 'config-change' in subs): subs.append('config-change') init = {'SUBSCRIBE': [{'name': x} for x in subs]} async for payload in self._ws(init=init): with suppress(KeyError): if reload_on_change and 'config-change' in payload and payload.get( 'config-change', [])['commit'] == 'ended': logger.debug( "Detected config change, refreshing config cache") await self.config() yield payload async def _ws(self, init, keepalive=True, timeout=30): pinger = None foo = {'UNSUBSCRIBE': []} foo.update(init) while True: try: ''' Make sure that before we launch the WebSocket we have a valid session id ''' while True: if await self.is_logged_in(): break logger.warning("Session died, trying a manual login.") await asyncio.sleep(5) await self.login() async with AsyncExitStack() as stack: ws = await stack.enter_async_context( self.session.ws_connect(f"{self.url}/ws/stats", headers=self.headers, origin=self.url, ssl=self.ssl)) pinger = await stack.enter_async_context( TaskEvery(ws_ping, ws, interval=30, sync_once=False)) foo.update({'SESSION_ID': self.session_id}) await ws.send_str(as_statd_string(foo)) data = '' while True: msg = await asyncio.wait_for(ws.receive(), timeout) if msg.type != WSMsgType.TEXT: logging.debug( f"got non text websocket data {msg.data!r} this probbaly means the socket was closed so let's start a fresh one" ) break data += msg.data temp1, temp2 = data.split('\n', 1) data_len = int(temp1) ''' If the payload is larger than the data len, process until it's less ''' while len(temp2) >= data_len: try: payload = json.loads(temp2[:data_len]) self.sysdata.update(payload) yield payload except Exception as e: logger.error(f"{e!r}") ''' Strip off the processed data and leave the next for another round ''' data = temp2[data_len:] if len(data) < 4: break temp1, temp2 = data.split('\n', 1) data_len = int(temp1) except asyncio.CancelledError as err: return except Exception as err: logger.debug(f"websocket loop raised {err!r}, ignoring") if not keepalive: return