refactor: use new fastapi lifespan instead of startup/shutdown events (#2294)

* refactor: use new fastapi lifespan instead of events
recommended use: https://fastapi.tiangolo.com/advanced/events/?h=lifespan
threw warnings in pytest
* make startup and shutdown functions
* nix: add override for asgi-lifespan

---------

Co-authored-by: Pavol Rusnak <pavol@rusnak.io>

flake.nix

@@ -33,6 +33,10 @@
             protobuf = prev.protobuf.override { preferWheel = true; };
             ruff = prev.ruff.override { preferWheel = true; };
             wallycore = prev.wallycore.override { preferWheel = true; };
+            # remove the following override when https://github.com/nix-community/poetry2nix/pull/1563 is merged
+            asgi-lifespan = prev.asgi-lifespan.overridePythonAttrs (
+              old: { buildInputs = (old.buildInputs or []) ++ [ prev.setuptools ]; }
+            );
           });
         };
       });

lnbits/app.py

@@ -7,6 +7,7 @@ import shutil
 import signal
 import sys
 import traceback
+from contextlib import asynccontextmanager
 from hashlib import sha256
 from http import HTTPStatus
 from pathlib import Path
@@ -68,6 +69,59 @@ from .tasks import (
 )
 
 
+async def startup(app: FastAPI):
+
+    # wait till migration is done
+    await migrate_databases()
+
+    # setup admin settings
+    await check_admin_settings()
+    await check_webpush_settings()
+
+    log_server_info()
+
+    # initialize WALLET
+    try:
+        set_wallet_class()
+    except Exception as e:
+        logger.error(f"Error initializing {settings.lnbits_backend_wallet_class}: {e}")
+        set_void_wallet_class()
+
+    # initialize funding source
+    await check_funding_source()
+
+    # register core routes
+    init_core_routers(app)
+
+    # check extensions after restart
+    if not settings.lnbits_extensions_deactivate_all:
+        await check_installed_extensions(app)
+        register_all_ext_routes(app)
+
+    if settings.lnbits_admin_ui:
+        initialize_server_logger()
+
+    # initialize tasks
+    register_async_tasks()
+
+
+async def shutdown():
+    # shutdown event
+    cancel_all_tasks()
+
+    # wait a bit to allow them to finish, so that cleanup can run without problems
+    await asyncio.sleep(0.1)
+    WALLET = get_wallet_class()
+    await WALLET.cleanup()
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    await startup(app)
+    yield
+    await shutdown()
+
+
 def create_app() -> FastAPI:
     configure_logger()
     app = FastAPI(
@@ -77,6 +131,7 @@ def create_app() -> FastAPI:
             "accounts system with plugins."
         ),
         version=settings.version,
+        lifespan=lifespan,
         license_info={
             "name": "MIT License",
             "url": "https://raw.githubusercontent.com/lnbits/lnbits/main/LICENSE",
@@ -117,10 +172,7 @@ def create_app() -> FastAPI:
     add_ip_block_middleware(app)
     add_ratelimit_middleware(app)
 
-    register_startup(app)
-    register_async_tasks(app)
     register_exception_handlers(app)
-    register_shutdown(app)
 
     return app
 
@@ -368,56 +420,6 @@ def register_all_ext_routes(app: FastAPI):
             logger.error(f"Could not load extension `{ext.code}`: {str(e)}")
 
 
-def register_startup(app: FastAPI):
-    @app.on_event("startup")
-    async def lnbits_startup():
-        try:
-            # wait till migration is done
-            await migrate_databases()
-
-            # setup admin settings
-            await check_admin_settings()
-            await check_webpush_settings()
-
-            log_server_info()
-
-            # initialize WALLET
-            try:
-                set_wallet_class()
-            except Exception as e:
-                logger.error(
-                    f"Error initializing {settings.lnbits_backend_wallet_class}: {e}"
-                )
-                set_void_wallet_class()
-
-            # initialize funding source
-            await check_funding_source()
-
-            init_core_routers(app)
-
-            # check extensions after restart
-            if not settings.lnbits_extensions_deactivate_all:
-                await check_installed_extensions(app)
-                register_all_ext_routes(app)
-
-            if settings.lnbits_admin_ui:
-                initialize_server_logger()
-
-        except Exception as e:
-            logger.error(str(e))
-            raise ImportError("Failed to run 'startup' event.")
-
-
-def register_shutdown(app: FastAPI):
-    @app.on_event("shutdown")
-    async def on_shutdown():
-        cancel_all_tasks()
-        # wait a bit to allow them to finish, so that cleanup can run without problems
-        await asyncio.sleep(0.1)
-        WALLET = get_wallet_class()
-        await WALLET.cleanup()
-
-
 def initialize_server_logger():
     super_user_hash = sha256(settings.super_user.encode("utf-8")).hexdigest()
 
@@ -465,9 +467,7 @@ def get_db_vendor_name():
     )
 
 
-def register_async_tasks(app):
+def register_async_tasks():
-    @app.on_event("startup")
-    async def listeners():
     create_permanent_task(check_pending_payments)
     create_permanent_task(invoice_listener)
     create_permanent_task(internal_invoice_listener)

poetry.lock

@@ -21,6 +21,20 @@ doc = ["Sphinx", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-
 test = ["anyio[trio]", "coverage[toml] (>=4.5)", "hypothesis (>=4.0)", "mock (>=4)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (>=0.17)"]
 trio = ["trio (<0.22)"]
 
+[[package]]
+name = "asgi-lifespan"
+version = "2.1.0"
+description = "Programmatic startup/shutdown of ASGI apps."
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "asgi-lifespan-2.1.0.tar.gz", hash = "sha256:5e2effaf0bfe39829cf2d64e7ecc47c7d86d676a6599f7afba378c31f5e3a308"},
+    {file = "asgi_lifespan-2.1.0-py3-none-any.whl", hash = "sha256:ed840706680e28428c01e14afb3875d7d76d3206f3d5b2f2294e059b5c23804f"},
+]
+
+[package.dependencies]
+sniffio = "*"
+
 [[package]]
 name = "asn1crypto"
 version = "1.5.1"
@@ -2934,4 +2948,4 @@ liquid = ["wallycore"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10 | ^3.9"
-content-hash = "fcc579d222f98204fbb9748cfd280a0f37a04cf5fc987dfccba02a66ce0f1f28"
+content-hash = "cbe93bb8afbda1cddb4e30721fb15a016b8fb1250d07ee06ff9365b8757c1710"

pyproject.toml

@@ -71,6 +71,7 @@ types-passlib = "^1.7.7.13"
 types-python-jose = "^3.3.4.8"
 openai = "^1.12.0"
 json5 = "^0.9.17"
+asgi-lifespan = "^2.1.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

tests/conftest.py

@@ -3,6 +3,7 @@ import asyncio
 from time import time
 
 import uvloop
+from asgi_lifespan import LifespanManager
 
 uvloop.install()
 
@@ -35,6 +36,7 @@ settings.lnbits_admin_extensions = []
 settings.lnbits_data_folder = "./tests/data"
 settings.lnbits_admin_ui = True
 settings.lnbits_extensions_default_install = []
+settings.lnbits_extensions_deactivate_all = True
 
 
 @pytest_asyncio.fixture(scope="session")
@@ -49,17 +51,16 @@ def event_loop():
 async def app():
     clean_database(settings)
     app = create_app()
-    await app.router.startup()
+    async with LifespanManager(app) as manager:
         settings.first_install = False
-    yield app
+        yield manager.app
-    await app.router.shutdown()
 
 
 @pytest_asyncio.fixture(scope="session")
 async def client(app):
-    client = AsyncClient(app=app, base_url=f"http://{settings.host}:{settings.port}")
+    url = f"http://{settings.host}:{settings.port}"
+    async with AsyncClient(app=app, base_url=url) as client:
         yield client
-    await client.aclose()
 
 
 @pytest.fixture(scope="session")