refactor: get_balance_delta and use pydantic model for openapi docs (#2492)

* refactor: `get_balance_delta` and use pydantic model for openapi docs

---------

Co-authored-by: Vlad Stan <stan.v.vlad@gmail.com>

lnbits/core/models.py

@@ -444,3 +444,12 @@ class WebPushSubscription(BaseModel):
     data: str
     host: str
     timestamp: str
+
+
+class BalanceDelta(BaseModel):
+    lnbits_balance_msats: int
+    node_balance_msats: int
+
+    @property
+    def delta_msats(self):
+        return self.node_balance_msats - self.lnbits_balance_msats

lnbits/core/services.py

@@ -59,7 +59,7 @@ from .crud import (
     update_super_user,
 )
 from .helpers import to_valid_user_id
-from .models import Payment, UserConfig, Wallet
+from .models import BalanceDelta, Payment, UserConfig, Wallet
 
 
 class PaymentError(Exception):
@@ -790,13 +790,14 @@ async def switch_to_voidwallet() -> None:
     settings.lnbits_backend_wallet_class = "VoidWallet"
 
 
-async def get_balance_delta() -> Tuple[int, int, int]:
+async def get_balance_delta() -> BalanceDelta:
     funding_source = get_funding_source()
-    total_balance = await get_total_balance()
+    status = await funding_source.status()
-    error_message, node_balance = await funding_source.status()
+    lnbits_balance = await get_total_balance()
-    if error_message:
+    return BalanceDelta(
-        raise Exception(error_message)
+        lnbits_balance_msats=lnbits_balance,
-    return node_balance - total_balance, node_balance, total_balance
+        node_balance_msats=status.balance_msat,
+    )
 
 
 def get_bolt11_expiry(invoice: Bolt11) -> datetime.datetime:

lnbits/core/tasks.py

@@ -63,7 +63,8 @@ async def watchdog_task():
             and funding_source.__class__.__name__ != "VoidWallet"
         ):
             try:
-                delta, *_ = await get_balance_delta()
+                balance = await get_balance_delta()
+                delta = balance.delta_msats
                 logger.debug(f"Running watchdog task. current delta: {delta}")
                 if delta + settings.lnbits_watchdog_delta <= 0:
                     logger.error(f"Switching to VoidWallet. current delta: {delta}")

lnbits/core/views/admin_api.py

@@ -8,7 +8,6 @@ from urllib.parse import urlparse
 
 from fastapi import APIRouter, Depends
 from fastapi.responses import FileResponse
-from starlette.exceptions import HTTPException
 
 from lnbits.core.models import User
 from lnbits.core.services import (
@@ -34,18 +33,7 @@ admin_router = APIRouter(tags=["Admin UI"], prefix="/admin")
     dependencies=[Depends(check_admin)],
 )
 async def api_auditor():
-    try:
+    return await get_balance_delta()
-        delta, node_balance, total_balance = await get_balance_delta()
-        return {
-            "delta_msats": int(delta),
-            "node_balance_msats": int(node_balance),
-            "lnbits_balance_msats": int(total_balance),
-        }
-    except Exception as exc:
-        raise HTTPException(
-            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
-            detail="Could not audit balance.",
-        ) from exc
 
 
 @admin_router.get(

tests/regtest/test_real_invoice.py

@@ -6,8 +6,7 @@ import pytest
 from lnbits import bolt11
 from lnbits.core.crud import get_standalone_payment, update_payment_details
 from lnbits.core.models import CreateInvoice, Payment
-from lnbits.core.services import fee_reserve_total
+from lnbits.core.services import fee_reserve_total, get_balance_delta
-from lnbits.core.views.admin_api import api_auditor
 from lnbits.core.views.payment_api import api_payment
 from lnbits.wallets import get_funding_source
 
@@ -21,8 +20,8 @@ from .helpers import (
 
 
 async def get_node_balance_sats():
-    audit = await api_auditor()
+    balance = await get_balance_delta()
-    return audit["node_balance_msats"] / 1000
+    return balance.node_balance_msats / 1000
 
 
 @pytest.mark.asyncio