From 99334845588304fb7bc289871fb4ab3bad296799 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?dni=20=E2=9A=A1?= <office@dnilabs.com>
Date: Mon, 13 May 2024 17:59:29 +0200
Subject: [PATCH] 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              |  9 +++++++++
 lnbits/core/services.py            | 15 ++++++++-------
 lnbits/core/tasks.py               |  3 ++-
 lnbits/core/views/admin_api.py     | 14 +-------------
 tests/regtest/test_real_invoice.py |  7 +++----
 5 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/lnbits/core/models.py b/lnbits/core/models.py
index 6a2f272c3..ef0a38e04 100644
--- a/lnbits/core/models.py
+++ b/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
diff --git a/lnbits/core/services.py b/lnbits/core/services.py
index c811f6099..2dacd6096 100644
--- a/lnbits/core/services.py
+++ b/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()
-    error_message, node_balance = await funding_source.status()
-    if error_message:
-        raise Exception(error_message)
-    return node_balance - total_balance, node_balance, total_balance
+    status = await funding_source.status()
+    lnbits_balance = await get_total_balance()
+    return BalanceDelta(
+        lnbits_balance_msats=lnbits_balance,
+        node_balance_msats=status.balance_msat,
+    )
 
 
 def get_bolt11_expiry(invoice: Bolt11) -> datetime.datetime:
diff --git a/lnbits/core/tasks.py b/lnbits/core/tasks.py
index e0c9c99a4..a229ff6bf 100644
--- a/lnbits/core/tasks.py
+++ b/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}")
diff --git a/lnbits/core/views/admin_api.py b/lnbits/core/views/admin_api.py
index eda01cc53..05d0b89ed 100644
--- a/lnbits/core/views/admin_api.py
+++ b/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:
-        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
+    return await get_balance_delta()
 
 
 @admin_router.get(
diff --git a/tests/regtest/test_real_invoice.py b/tests/regtest/test_real_invoice.py
index bdf043b0f..926989055 100644
--- a/tests/regtest/test_real_invoice.py
+++ b/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.views.admin_api import api_auditor
+from lnbits.core.services import fee_reserve_total, get_balance_delta
 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()
-    return audit["node_balance_msats"] / 1000
+    balance = await get_balance_delta()
+    return balance.node_balance_msats / 1000
 
 
 @pytest.mark.asyncio