From cae72b5848f088a118e64d828783cd43a9e3efd4 Mon Sep 17 00:00:00 2001
From: carla <kirkcohenc@gmail.com>
Date: Mon, 7 Jun 2021 10:56:34 +0200
Subject: [PATCH] loopin/test: move invoice updates to loopin test context

---
 loopin_test.go             | 37 +++++++++----------------------------
 loopin_testcontext_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 28 deletions(-)

diff --git a/loopin_test.go b/loopin_test.go
index 8189558..93b6454 100644
--- a/loopin_test.go
+++ b/loopin_test.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"testing"
 
-	"github.com/lightninglabs/lndclient"
 	"github.com/lightninglabs/loop/loopdb"
 	"github.com/lightninglabs/loop/swap"
 	"github.com/lightninglabs/loop/test"
@@ -89,17 +88,11 @@ func TestLoopInSuccess(t *testing.T) {
 	<-ctx.lnd.RegisterSpendChannel
 
 	// Client starts listening for swap invoice updates.
-	subscription := <-ctx.lnd.SingleInvoiceSubcribeChannel
-	if subscription.Hash != ctx.server.swapHash {
-		t.Fatal("client subscribing to wrong invoice")
-	}
+	ctx.assertSubscribeInvoice(ctx.server.swapHash)
 
 	// Server has already paid invoice before spending the htlc. Signal
 	// settled.
-	subscription.Update <- lndclient.InvoiceUpdate{
-		State:   channeldb.ContractSettled,
-		AmtPaid: 49000,
-	}
+	ctx.updateInvoiceState(49000, channeldb.ContractSettled)
 
 	// Swap is expected to move to the state InvoiceSettled
 	ctx.assertState(loopdb.StateInvoiceSettled)
@@ -257,10 +250,7 @@ func testLoopInTimeout(t *testing.T,
 	<-ctx.lnd.RegisterSpendChannel
 
 	// Client starts listening for swap invoice updates.
-	subscription := <-ctx.lnd.SingleInvoiceSubcribeChannel
-	if subscription.Hash != ctx.server.swapHash {
-		t.Fatal("client subscribing to wrong invoice")
-	}
+	ctx.assertSubscribeInvoice(ctx.server.swapHash)
 
 	// Let htlc expire.
 	ctx.blockEpochChan <- s.LoopInContract.CltvExpiry
@@ -294,10 +284,8 @@ func testLoopInTimeout(t *testing.T,
 	// safely cancel the swap invoice.
 	<-ctx.lnd.FailInvoiceChannel
 
-	// Signal the the invoice was canceled.
-	subscription.Update <- lndclient.InvoiceUpdate{
-		State: channeldb.ContractCanceled,
-	}
+	// Signal that the invoice was canceled.
+	ctx.updateInvoiceState(0, channeldb.ContractCanceled)
 
 	ctx.assertState(loopdb.StateFailTimeout)
 	state := ctx.store.assertLoopInState(loopdb.StateFailTimeout)
@@ -501,18 +489,12 @@ func testLoopInResume(t *testing.T, state loopdb.SwapState, expired bool,
 	<-ctx.lnd.RegisterSpendChannel
 
 	// Client starts listening for swap invoice updates.
-	subscription := <-ctx.lnd.SingleInvoiceSubcribeChannel
-	if subscription.Hash != testPreimage.Hash() {
-		t.Fatal("client subscribing to wrong invoice")
-	}
+	ctx.assertSubscribeInvoice(testPreimage.Hash())
 
 	// Server has already paid invoice before spending the htlc. Signal
 	// settled.
-	invoiceUpdate := lndclient.InvoiceUpdate{
-		State:   channeldb.ContractSettled,
-		AmtPaid: 49000,
-	}
-	subscription.Update <- invoiceUpdate
+	amtPaid := btcutil.Amount(49000)
+	ctx.updateInvoiceState(amtPaid, channeldb.ContractSettled)
 
 	// Swap is expected to move to the state InvoiceSettled
 	ctx.assertState(loopdb.StateInvoiceSettled)
@@ -537,7 +519,6 @@ func testLoopInResume(t *testing.T, state loopdb.SwapState, expired bool,
 	// We expect our server fee to reflect as the difference between htlc
 	// value and invoice amount paid. We use our original on-chain cost, set
 	// earlier in the test, because we expect this value to be unchanged.
-	cost.Server = btcutil.Amount(htlcTx.TxOut[0].Value) -
-		invoiceUpdate.AmtPaid
+	cost.Server = btcutil.Amount(htlcTx.TxOut[0].Value) - amtPaid
 	require.Equal(t, cost, finalState.Cost)
 }
diff --git a/loopin_testcontext_test.go b/loopin_testcontext_test.go
index 1af542d..7f0c0a8 100644
--- a/loopin_testcontext_test.go
+++ b/loopin_testcontext_test.go
@@ -4,9 +4,14 @@ import (
 	"testing"
 	"time"
 
+	"github.com/btcsuite/btcutil"
+	"github.com/lightninglabs/lndclient"
 	"github.com/lightninglabs/loop/loopdb"
 	"github.com/lightninglabs/loop/sweep"
 	"github.com/lightninglabs/loop/test"
+	"github.com/lightningnetwork/lnd/channeldb"
+	"github.com/lightningnetwork/lnd/lntypes"
+	"github.com/stretchr/testify/require"
 )
 
 type loopInTestContext struct {
@@ -18,6 +23,8 @@ type loopInTestContext struct {
 	cfg            *executeConfig
 	statusChan     chan SwapInfo
 	blockEpochChan chan interface{}
+
+	swapInvoiceSubscription *test.SingleInvoiceSubscription
 }
 
 func newLoopInTestContext(t *testing.T) *loopInTestContext {
@@ -61,3 +68,20 @@ func (c *loopInTestContext) assertState(expectedState loopdb.SwapState) {
 			state.State)
 	}
 }
+
+// assertSubscribeInvoice asserts that the client subscribes to invoice updates
+// for our swap invoice.
+func (c *loopInTestContext) assertSubscribeInvoice(hash lntypes.Hash) {
+	c.swapInvoiceSubscription = <-c.lnd.SingleInvoiceSubcribeChannel
+	require.Equal(c.t, hash, c.swapInvoiceSubscription.Hash)
+}
+
+// updateInvoiceState mocks an update to our swap invoice state.
+func (c *loopInTestContext) updateInvoiceState(amount btcutil.Amount,
+	state channeldb.ContractState) {
+
+	c.swapInvoiceSubscription.Update <- lndclient.InvoiceUpdate{
+		AmtPaid: amount,
+		State:   state,
+	}
+}