From 09029bfdec15e1b9e2e1a56329ad9e789c556394 Mon Sep 17 00:00:00 2001
From: Wilmer Paulino <wilmer.paulino@gmail.com>
Date: Tue, 1 Oct 2019 11:21:17 -0400
Subject: [PATCH] test: allow custom fee estimates

---
 test/lnd_services_mock.go | 16 ++++++++++++----
 test/walletkit_mock.go    | 16 +++++++++++++---
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/test/lnd_services_mock.go b/test/lnd_services_mock.go
index b085c4f..c7a7fa7 100644
--- a/test/lnd_services_mock.go
+++ b/test/lnd_services_mock.go
@@ -6,12 +6,12 @@ import (
 	"time"
 
 	"github.com/btcsuite/btcd/chaincfg"
-	"github.com/lightningnetwork/lnd/lntypes"
-	"github.com/lightningnetwork/lnd/zpay32"
-
 	"github.com/btcsuite/btcd/wire"
 	"github.com/lightninglabs/loop/lndclient"
 	"github.com/lightningnetwork/lnd/chainntnfs"
+	"github.com/lightningnetwork/lnd/lntypes"
+	"github.com/lightningnetwork/lnd/lnwallet"
+	"github.com/lightningnetwork/lnd/zpay32"
 )
 
 var testStartingHeight = int32(600)
@@ -20,7 +20,9 @@ var testStartingHeight = int32(600)
 // tests.
 func NewMockLnd() *LndMockServices {
 	lightningClient := &mockLightningClient{}
-	walletKit := &mockWalletKit{}
+	walletKit := &mockWalletKit{
+		feeEstimates: make(map[int32]lnwallet.SatPerKWeight),
+	}
 	chainNotifier := &mockChainNotifier{}
 	signer := &mockSigner{}
 	invoices := &mockInvoices{}
@@ -197,3 +199,9 @@ func (s *LndMockServices) DecodeInvoice(request string) (*zpay32.Invoice,
 
 	return zpay32.Decode(request, s.ChainParams)
 }
+
+func (s *LndMockServices) SetFeeEstimate(confTarget int32,
+	feeEstimate lnwallet.SatPerKWeight) {
+
+	s.WalletKit.(*mockWalletKit).feeEstimates[confTarget] = feeEstimate
+}
diff --git a/test/walletkit_mock.go b/test/walletkit_mock.go
index c8dc5bb..8ebc57f 100644
--- a/test/walletkit_mock.go
+++ b/test/walletkit_mock.go
@@ -8,15 +8,19 @@ import (
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btcutil"
+	"github.com/lightninglabs/loop/lndclient"
 	"github.com/lightningnetwork/lnd/keychain"
 	"github.com/lightningnetwork/lnd/lnwallet"
 )
 
 type mockWalletKit struct {
-	lnd      *LndMockServices
-	keyIndex int32
+	lnd          *LndMockServices
+	keyIndex     int32
+	feeEstimates map[int32]lnwallet.SatPerKWeight
 }
 
+var _ lndclient.WalletKitClient = (*mockWalletKit)(nil)
+
 func (m *mockWalletKit) DeriveNextKey(ctx context.Context, family int32) (
 	*keychain.KeyDescriptor, error) {
 
@@ -87,9 +91,15 @@ func (m *mockWalletKit) SendOutputs(ctx context.Context, outputs []*wire.TxOut,
 
 func (m *mockWalletKit) EstimateFee(ctx context.Context, confTarget int32) (
 	lnwallet.SatPerKWeight, error) {
+
 	if confTarget <= 1 {
 		return 0, errors.New("conf target must be greater than 1")
 	}
 
-	return 10000, nil
+	feeEstimate, ok := m.feeEstimates[confTarget]
+	if !ok {
+		return 10000, nil
+	}
+
+	return feeEstimate, nil
 }