From 033fd3c83d1135070b231ea618ff2076a51c9d7e Mon Sep 17 00:00:00 2001
From: Vadym Popov <me@vpopov.org>
Date: Sun, 12 Aug 2018 16:20:57 +0300
Subject: [PATCH] htlcswitch: add test for integrated control tower

---
 htlcswitch/link_test.go   |  4 ++--
 htlcswitch/mock.go        | 26 +++++++++++++++++++-------
 htlcswitch/switch_test.go | 31 +++++--------------------------
 3 files changed, 26 insertions(+), 35 deletions(-)

diff --git a/htlcswitch/link_test.go b/htlcswitch/link_test.go
index a5f417be..b48587e6 100644
--- a/htlcswitch/link_test.go
+++ b/htlcswitch/link_test.go
@@ -3755,8 +3755,8 @@ func TestChannelLinkAcceptDuplicatePayment(t *testing.T) {
 		n.firstBobChannelLink.ShortChanID(), htlc,
 		newMockDeobfuscator(),
 	)
-	if err != nil {
-		t.Fatalf("error shouldn't have been received got: %v", err)
+	if err != ErrAlreadyPaid {
+		t.Fatalf("ErrAlreadyPaid should have been received got: %v", err)
 	}
 }
 
diff --git a/htlcswitch/mock.go b/htlcswitch/mock.go
index 9df91993..cd563ce7 100644
--- a/htlcswitch/mock.go
+++ b/htlcswitch/mock.go
@@ -124,14 +124,26 @@ type mockServer struct {
 
 var _ lnpeer.Peer = (*mockServer)(nil)
 
-func initSwitchWithDB(startingHeight uint32, db *channeldb.DB) (*Switch, error) {
-	if db == nil {
-		tempPath, err := ioutil.TempDir("", "switchdb")
-		if err != nil {
-			return nil, err
-		}
+func initDB() (*channeldb.DB, error) {
+	tempPath, err := ioutil.TempDir("", "switchdb")
+	if err != nil {
+		return nil, err
+	}
 
-		db, err = channeldb.Open(tempPath)
+	db, err := channeldb.Open(tempPath)
+	if err != nil {
+		return nil, err
+	}
+
+	return db, err
+}
+
+
+func initSwitchWithDB(startingHeight uint32, db *channeldb.DB) (*Switch, error) {
+	var err error
+
+	if db == nil {
+		db, err = initDB()
 		if err != nil {
 			return nil, err
 		}
diff --git a/htlcswitch/switch_test.go b/htlcswitch/switch_test.go
index 7610a797..6dd8980f 100644
--- a/htlcswitch/switch_test.go
+++ b/htlcswitch/switch_test.go
@@ -1678,7 +1678,9 @@ func TestSwitchSendPayment(t *testing.T) {
 		}
 
 	case err := <-errChan:
-		t.Fatalf("unable to send payment: %v", err)
+		if err != ErrPaymentInFlight {
+			t.Fatalf("unable to send payment: %v", err)
+		}
 	case <-time.After(time.Second):
 		t.Fatal("request was not propagated to destination")
 	}
@@ -1695,11 +1697,11 @@ func TestSwitchSendPayment(t *testing.T) {
 		t.Fatal("request was not propagated to destination")
 	}
 
-	if s.numPendingPayments() != 2 {
+	if s.numPendingPayments() != 1 {
 		t.Fatal("wrong amount of pending payments")
 	}
 
-	if s.circuits.NumOpen() != 2 {
+	if s.circuits.NumOpen() != 1 {
 		t.Fatal("wrong amount of circuits")
 	}
 
@@ -1735,29 +1737,6 @@ func TestSwitchSendPayment(t *testing.T) {
 		t.Fatal("err wasn't received")
 	}
 
-	packet = &htlcPacket{
-		outgoingChanID: aliceChannelLink.ShortChanID(),
-		outgoingHTLCID: 1,
-		htlc: &lnwire.UpdateFailHTLC{
-			Reason: reason,
-		},
-	}
-
-	// Send second failure response and check that user were able to
-	// receive the error.
-	if err := s.forward(packet); err != nil {
-		t.Fatalf("can't forward htlc packet: %v", err)
-	}
-
-	select {
-	case err := <-errChan:
-		if err.Error() != errors.New(lnwire.CodeIncorrectPaymentAmount).Error() {
-			t.Fatal("err wasn't received")
-		}
-	case <-time.After(time.Second):
-		t.Fatal("err wasn't received")
-	}
-
 	if s.numPendingPayments() != 0 {
 		t.Fatal("wrong amount of pending payments")
 	}