diff --git a/lnbits/core/views/payment_api.py b/lnbits/core/views/payment_api.py index f9bcb4de2..648237ffc 100644 --- a/lnbits/core/views/payment_api.py +++ b/lnbits/core/views/payment_api.py @@ -1,5 +1,6 @@ from hashlib import sha256 from http import HTTPStatus +from secrets import token_hex from fastapi import ( APIRouter, @@ -37,8 +38,9 @@ from lnbits.core.models import ( UpdatePaymentExtra, ) from lnbits.core.models.payments import UpdatePaymentLabels -from lnbits.core.models.users import AccountId +from lnbits.core.models.users import AccountId, UserLabel from lnbits.core.models.wallets import BaseWalletTypeInfo +from lnbits.core.services.users import update_user_account from lnbits.db import Filters, Page from lnbits.decorators import ( WalletTypeInfo, @@ -51,6 +53,7 @@ from lnbits.decorators import ( from lnbits.helpers import ( filter_dict_keys, generate_filter_params_openapi, + is_valid_label, ) from lnbits.wallets.base import InvoiceResponse @@ -291,9 +294,24 @@ async def api_update_payment_labels( if not account: raise HTTPException(HTTPStatus.NOT_FOUND, "Account does not exist.") - # only keep labels that belong to the user user_label_names = [label.name for label in account.extra.labels] - payment.labels = [label for label in data.labels if label in user_label_names] + updated_account = False + for label_name in data.labels: + if label_name not in user_label_names: + if not is_valid_label(label_name): + raise HTTPException( + HTTPStatus.BAD_REQUEST, f"Invalid label name: '{label_name}'." + ) + account.extra.labels.append( + UserLabel(name=label_name, color=f"#{token_hex(3)}") + ) + user_label_names.append(label_name) + updated_account = True + + if updated_account: + await update_user_account(account) + + payment.labels = data.labels await update_payment(payment) return SimpleStatus(success=True, message="Payment labels updated.") diff --git a/tests/api/test_payment_api.py b/tests/api/test_payment_api.py index 95339b2d0..447feaffb 100644 --- a/tests/api/test_payment_api.py +++ b/tests/api/test_payment_api.py @@ -376,6 +376,62 @@ async def test_payment_extra_update_requires_successful_payment( ) +@pytest.mark.anyio +async def test_api_update_payment_labels( + client, + to_wallet, + adminkey_headers_to, +): + payment_hash = uuid4().hex + checking_id = await _create_payment( + to_wallet.id, + amount_msat=1_000, + payment_hash=payment_hash, + status=PaymentState.SUCCESS, + ) + + # 1. Update labels with new valid labels + response = await client.put( + f"/api/v1/payments/{payment_hash}/labels", + headers=adminkey_headers_to, + json={"labels": ["income", "restaurant"]}, + ) + assert response.status_code == 200 + assert response.json()["success"] is True + + # 2. Check that the labels were updated on the payment + payment = await get_payment(checking_id) + assert payment.labels == ["income", "restaurant"] + + # 3. Check that the new labels were auto-created in the user's account config + from lnbits.core.crud.users import get_account + + account = await get_account(to_wallet.user) + assert account is not None + user_labels = {label.name: label.color for label in account.extra.labels} + assert "income" in user_labels + assert "restaurant" in user_labels + assert ( + user_labels["income"] is not None + and user_labels["income"].startswith("#") + and len(user_labels["income"]) == 7 + ) + assert ( + user_labels["restaurant"] is not None + and user_labels["restaurant"].startswith("#") + and len(user_labels["restaurant"]) == 7 + ) + + # 4. Check that invalid labels are rejected + response = await client.put( + f"/api/v1/payments/{payment_hash}/labels", + headers=adminkey_headers_to, + json={"labels": ["invalid!label"]}, + ) + assert response.status_code == 400 + assert "Invalid label name" in response.json()["detail"] + + async def _create_payment( wallet_id: str, *,