diff --git a/fluid/tools_fastapi/backdoor.py b/fluid/tools_fastapi/backdoor.py index 2d59bbc..ee4cfcf 100644 --- a/fluid/tools_fastapi/backdoor.py +++ b/fluid/tools_fastapi/backdoor.py @@ -1,11 +1,12 @@ from functools import partial +from typing import Any from fastapi import FastAPI from fluid.utils.backdoor import ConsoleManager -def setup(app: FastAPI) -> None: - console = ConsoleManager() +def setup(app: FastAPI, **kwargs: Any) -> None: + console = ConsoleManager(**kwargs) app.add_event_handler("startup", partial(console.on_startup, app)) app.add_event_handler("shutdown", partial(console.on_cleanup, app)) diff --git a/fluid/utils/backdoor.py b/fluid/utils/backdoor.py index 591a031..4393ffa 100644 --- a/fluid/utils/backdoor.py +++ b/fluid/utils/backdoor.py @@ -1,5 +1,6 @@ import logging import sys +from dataclasses import dataclass from functools import partial import aioconsole @@ -22,16 +23,11 @@ def get_default_banner(self): return f"Python {sys.version} on {sys.platform}\n{cprt}\n{CONSOLE_MESSAGE}" +@dataclass class ConsoleManager: aio_console = None - - def __init__( - self, - port: int = settings.AIO_BACKDOOR_PORT, - host: str = "0.0.0.0", - ) -> None: - self.port = port - self.host = host + port: int = settings.AIO_BACKDOOR_PORT + host: str = "0.0.0.0" async def on_startup(self, app) -> None: self.aio_console = await aioconsole.start_interactive_server( diff --git a/tests/scheduler/conftest.py b/tests/scheduler/conftest.py index 44152df..610f943 100644 --- a/tests/scheduler/conftest.py +++ b/tests/scheduler/conftest.py @@ -9,11 +9,13 @@ from fluid.scheduler import TaskManager, TaskScheduler from fluid.scheduler.broker import RedisTaskBroker from fluid.scheduler.endpoints import get_task_manger, setup_fastapi +from fluid.tools_fastapi import backdoor from tests.scheduler.tasks import TaskClient, task_application @asynccontextmanager async def start_fastapi(app: FastAPI) -> AsyncIterator: + backdoor.setup(app, port=0) async with app.router.lifespan_context(app): yield app