From e557374183aac8497bc909885200e6357e419f4a Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Sat, 4 Jan 2025 16:12:04 +0200
Subject: [PATCH] accounts: use UpsertAccountPayment in removePayment

Instead of UpdateAccount.
---
 accounts/errors.go     |  4 ++++
 accounts/interface.go  | 11 +++++++++++
 accounts/service.go    | 25 +++++++++++++------------
 accounts/store_kvdb.go |  2 ++
 accounts/store_test.go |  9 +++++++++
 5 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/accounts/errors.go b/accounts/errors.go
index 241060a2a..3fe443460 100644
--- a/accounts/errors.go
+++ b/accounts/errors.go
@@ -13,4 +13,8 @@ var (
 	// if the WithErrAlreadySucceeded option is used and the payment has
 	// already succeeded.
 	ErrAlreadySucceeded = errors.New("payment has already succeeded")
+
+	// ErrPaymentUnknown is returned by the UpsertAccountPayment method if
+	// the WithErrIfUnknown option is used and the payment is not yet known.
+	ErrPaymentUnknown = errors.New("payment unknown")
 )
diff --git a/accounts/interface.go b/accounts/interface.go
index b2e11b3e0..bd75d9be0 100644
--- a/accounts/interface.go
+++ b/accounts/interface.go
@@ -329,6 +329,7 @@ type upsertAcctPaymentOption struct {
 	errIfAlreadyPending   bool
 	usePendingAmount      bool
 	errIfAlreadySucceeded bool
+	errIfUnknown          bool
 }
 
 // newUpsertPaymentOption creates a new upsertAcctPaymentOption with default
@@ -339,6 +340,7 @@ func newUpsertPaymentOption() *upsertAcctPaymentOption {
 		errIfAlreadyPending:   false,
 		usePendingAmount:      false,
 		errIfAlreadySucceeded: false,
+		errIfUnknown:          false,
 	}
 }
 
@@ -378,3 +380,12 @@ func WithPendingAmount() UpsertPaymentOption {
 		o.usePendingAmount = true
 	}
 }
+
+// WithErrIfUnknown is a functional option that can be passed to the
+// UpsertAccountPayment method to indicate that the ErrPaymentUnknown error
+// should be returned if the payment is not associated with the account.
+func WithErrIfUnknown() UpsertPaymentOption {
+	return func(o *upsertAcctPaymentOption) {
+		o.errIfUnknown = true
+	}
+}
diff --git a/accounts/service.go b/accounts/service.go
index de516f5a6..b8fee1c10 100644
--- a/accounts/service.go
+++ b/accounts/service.go
@@ -845,23 +845,24 @@ func (s *InterceptorService) removePayment(ctx context.Context,
 		return nil
 	}
 
-	account, err := s.store.Account(ctx, pendingPayment.accountID)
-	if err != nil {
-		return err
-	}
-
 	pendingPayment.cancel()
 	delete(s.pendingPayments, hash)
 
-	// Have we associated the payment with the account already?
-	_, ok = account.Payments[hash]
-	if !ok {
-		return nil
+	_, err := s.store.UpsertAccountPayment(
+		ctx, pendingPayment.accountID, hash, 0, status,
+		// We don't want the payment to be inserted if it isn't already
+		// known. So we pass in this option to ensure that the call
+		// exists early if the payment is unknown.
+		WithErrIfUnknown(),
+		// Otherwise, we just want to update the status of the payment
+		// and use the existing pending amount.
+		WithPendingAmount(),
+	)
+	if err != nil && !errors.Is(err, ErrPaymentUnknown) {
+		return fmt.Errorf("error updating account: %w", err)
 	}
 
-	// If we did, let's set the status correctly in the DB now.
-	account.Payments[hash].Status = status
-	return s.store.UpdateAccount(ctx, account)
+	return nil
 }
 
 // successState returns true if a payment was completed successfully.
diff --git a/accounts/store_kvdb.go b/accounts/store_kvdb.go
index 18115e75e..76f8e8423 100644
--- a/accounts/store_kvdb.go
+++ b/accounts/store_kvdb.go
@@ -240,6 +240,8 @@ func (s *BoltStore) UpsertAccountPayment(_ context.Context, id AccountID,
 			if opts.usePendingAmount {
 				fullAmount = entry.FullAmount
 			}
+		} else if opts.errIfUnknown {
+			return ErrPaymentUnknown
 		}
 
 		account.Payments[paymentHash] = &PaymentEntry{
diff --git a/accounts/store_test.go b/accounts/store_test.go
index a7f35fcba..ed8dc9e8f 100644
--- a/accounts/store_test.go
+++ b/accounts/store_test.go
@@ -275,6 +275,15 @@ func TestAccountUpdateMethods(t *testing.T) {
 		// error.
 		err = store.DeleteAccountPayment(ctx, acct.ID, hash1)
 		require.ErrorContains(t, err, "is not associated")
+
+		// Try once more to insert a payment that is currently unknown
+		// but this time add the WithErrIfUnknown option. This should
+		// return the ErrPaymentUnknown error.
+		_, err = store.UpsertAccountPayment(
+			ctx, acct.ID, hash1, 600, lnrpc.Payment_SUCCEEDED,
+			WithErrIfUnknown(),
+		)
+		require.ErrorIs(t, err, ErrPaymentUnknown)
 	})
 }