From e0c52e447339311c2a1e09825c9a48bd7fa2c71e Mon Sep 17 00:00:00 2001
From: carla <kirkcohenc@gmail.com>
Date: Fri, 23 Apr 2021 08:39:40 +0200
Subject: [PATCH] routing/test: close payment result channel on shutdown,
 mimicking switch

This commit updates our mock to more closely follow the behavior of the
switch for mocked calls to GetPaymentResult. As it stands, our tests
send a test-created error from the switch when we want to mock shutdown.
In reality, the switch will close its result channel, so we update this
test to follow that behavior. This matters for the commit that follows,
because we start checking the error our payments return. If we have an
error from the switch, our tests will fail with an error that we do
not encounter in practice.
---
 routing/mock_test.go              | 19 +++++++++++--------
 routing/payment_lifecycle_test.go | 26 +++++++++-----------------
 2 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/routing/mock_test.go b/routing/mock_test.go
index ae67075d9..2caf9336c 100644
--- a/routing/mock_test.go
+++ b/routing/mock_test.go
@@ -155,10 +155,9 @@ func (m *mockPaymentSession) RequestRoute(_, _ lnwire.MilliSatoshi,
 }
 
 type mockPayer struct {
-	sendResult       chan error
-	paymentResultErr chan error
-	paymentResult    chan *htlcswitch.PaymentResult
-	quit             chan struct{}
+	sendResult    chan error
+	paymentResult chan *htlcswitch.PaymentResult
+	quit          chan struct{}
 }
 
 var _ PaymentAttemptDispatcher = (*mockPayer)(nil)
@@ -180,12 +179,16 @@ func (m *mockPayer) GetPaymentResult(paymentID uint64, _ lntypes.Hash,
 	_ htlcswitch.ErrorDecrypter) (<-chan *htlcswitch.PaymentResult, error) {
 
 	select {
-	case res := <-m.paymentResult:
+	case res, ok := <-m.paymentResult:
 		resChan := make(chan *htlcswitch.PaymentResult, 1)
-		resChan <- res
+		if !ok {
+			close(resChan)
+		} else {
+			resChan <- res
+		}
+
 		return resChan, nil
-	case err := <-m.paymentResultErr:
-		return nil, err
+
 	case <-m.quit:
 		return nil, fmt.Errorf("test quitting")
 	}
diff --git a/routing/payment_lifecycle_test.go b/routing/payment_lifecycle_test.go
index 62c85c36c..a122f8704 100644
--- a/routing/payment_lifecycle_test.go
+++ b/routing/payment_lifecycle_test.go
@@ -2,7 +2,6 @@ package routing
 
 import (
 	"crypto/rand"
-	"fmt"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -608,7 +607,7 @@ func testPaymentLifecycle(t *testing.T, test paymentLifecycleTestCase,
 	// setupRouter is a helper method that creates and starts the router in
 	// the desired configuration for this test.
 	setupRouter := func() (*ChannelRouter, chan error,
-		chan *htlcswitch.PaymentResult, chan error) {
+		chan *htlcswitch.PaymentResult) {
 
 		chain := newMockChain(startingBlockHeight)
 		chainView := newMockChainView(chain)
@@ -616,13 +615,11 @@ func testPaymentLifecycle(t *testing.T, test paymentLifecycleTestCase,
 		// We set uo the use the following channels and a mock Payer to
 		// synchonize with the interaction to the Switch.
 		sendResult := make(chan error)
-		paymentResultErr := make(chan error)
 		paymentResult := make(chan *htlcswitch.PaymentResult)
 
 		payer := &mockPayer{
-			sendResult:       sendResult,
-			paymentResult:    paymentResult,
-			paymentResultErr: paymentResultErr,
+			sendResult:    sendResult,
+			paymentResult: paymentResult,
 		}
 
 		router, err := New(Config{
@@ -675,10 +672,10 @@ func testPaymentLifecycle(t *testing.T, test paymentLifecycleTestCase,
 			t.Fatalf("did not fetch in flight payments at startup")
 		}
 
-		return router, sendResult, paymentResult, paymentResultErr
+		return router, sendResult, paymentResult
 	}
 
-	router, sendResult, getPaymentResult, getPaymentResultErr := setupRouter()
+	router, sendResult, getPaymentResult := setupRouter()
 	defer func() {
 		if err := router.Stop(); err != nil {
 			t.Fatal(err)
@@ -859,13 +856,9 @@ func testPaymentLifecycle(t *testing.T, test paymentLifecycleTestCase,
 
 		// In this step we manually stop the router.
 		case stopRouter:
-			select {
-			case getPaymentResultErr <- fmt.Errorf(
-				"shutting down"):
-			case <-time.After(stepTimeout):
-				t.Fatalf("unable to send payment " +
-					"result error")
-			}
+			// On shutdown, the switch closes our result channel.
+			// Mimic this behavior in our mock.
+			close(getPaymentResult)
 
 			if err := router.Stop(); err != nil {
 				t.Fatalf("unable to restart: %v", err)
@@ -873,8 +866,7 @@ func testPaymentLifecycle(t *testing.T, test paymentLifecycleTestCase,
 
 		// In this step we manually start the router.
 		case startRouter:
-			router, sendResult, getPaymentResult,
-				getPaymentResultErr = setupRouter()
+			router, sendResult, getPaymentResult = setupRouter()
 
 		// In this state we expect to receive an error for the
 		// original payment made.