From ab75cd3a4dcd7d2aeff04dbe59bd7d8c25284756 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Tue, 16 Jan 2018 21:23:56 +0100
Subject: [PATCH] htlcswitch/test_utils: set BatchTicker mock and BatchSize

This commit adds the mockTicker and BatchSize to the link
config. It also exits the goroutines draining the HtlcUpdates
gracefully.
---
 htlcswitch/test_utils.go | 49 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 45 insertions(+), 4 deletions(-)

diff --git a/htlcswitch/test_utils.go b/htlcswitch/test_utils.go
index 85718582..ca9609c7 100644
--- a/htlcswitch/test_utils.go
+++ b/htlcswitch/test_utils.go
@@ -437,17 +437,21 @@ type threeHopNetwork struct {
 	aliceServer      *mockServer
 	aliceChannelLink *channelLink
 	aliceBlockEpoch  chan *chainntnfs.BlockEpoch
+	aliceTicker      *time.Ticker
 
 	firstBobChannelLink *channelLink
 	bobFirstBlockEpoch  chan *chainntnfs.BlockEpoch
+	firstBobTicker      *time.Ticker
 
 	bobServer            *mockServer
 	secondBobChannelLink *channelLink
 	bobSecondBlockEpoch  chan *chainntnfs.BlockEpoch
+	secondBobTicker      *time.Ticker
 
 	carolChannelLink *channelLink
 	carolServer      *mockServer
 	carolBlockEpoch  chan *chainntnfs.BlockEpoch
+	carolTicker      *time.Ticker
 
 	feeEstimator *mockFeeEstimator
 
@@ -625,6 +629,11 @@ func (n *threeHopNetwork) stop() {
 		done <- struct{}{}
 	}()
 
+	n.aliceTicker.Stop()
+	n.firstBobTicker.Stop()
+	n.secondBobTicker.Stop()
+	n.carolTicker.Stop()
+
 	for i := 0; i < 3; i++ {
 		<-done
 	}
@@ -743,6 +752,7 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 		Cancel: func() {
 		},
 	}
+	aliceTicker := time.NewTicker(50 * time.Millisecond)
 	aliceChannelLink := NewChannelLink(
 		ChannelLinkConfig{
 			FwrdingPolicy:     globalPolicy,
@@ -763,6 +773,8 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 			},
 			ChainEvents: &contractcourt.ChainEventSubscription{},
 			SyncStates:  true,
+			BatchTicker: &mockTicker{aliceTicker.C},
+			BatchSize:   10,
 		},
 		aliceChannel,
 		startingHeight,
@@ -772,7 +784,11 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 	}
 	go func() {
 		for {
-			<-aliceChannelLink.(*channelLink).htlcUpdates
+			select {
+			case <-aliceChannelLink.(*channelLink).htlcUpdates:
+			case <-aliceChannelLink.(*channelLink).quit:
+				return
+			}
 		}
 	}()
 
@@ -782,6 +798,7 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 		Cancel: func() {
 		},
 	}
+	firstBobTicker := time.NewTicker(50 * time.Millisecond)
 	firstBobChannelLink := NewChannelLink(
 		ChannelLinkConfig{
 			FwrdingPolicy:     globalPolicy,
@@ -802,6 +819,8 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 			},
 			ChainEvents: &contractcourt.ChainEventSubscription{},
 			SyncStates:  true,
+			BatchTicker: &mockTicker{firstBobTicker.C},
+			BatchSize:   10,
 		},
 		firstBobChannel,
 		startingHeight,
@@ -811,7 +830,11 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 	}
 	go func() {
 		for {
-			<-firstBobChannelLink.(*channelLink).htlcUpdates
+			select {
+			case <-firstBobChannelLink.(*channelLink).htlcUpdates:
+			case <-firstBobChannelLink.(*channelLink).quit:
+				return
+			}
 		}
 	}()
 
@@ -821,6 +844,7 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 		Cancel: func() {
 		},
 	}
+	secondBobTicker := time.NewTicker(50 * time.Millisecond)
 	secondBobChannelLink := NewChannelLink(
 		ChannelLinkConfig{
 			FwrdingPolicy:     globalPolicy,
@@ -841,6 +865,8 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 			},
 			ChainEvents: &contractcourt.ChainEventSubscription{},
 			SyncStates:  true,
+			BatchTicker: &mockTicker{secondBobTicker.C},
+			BatchSize:   10,
 		},
 		secondBobChannel,
 		startingHeight,
@@ -850,7 +876,11 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 	}
 	go func() {
 		for {
-			<-secondBobChannelLink.(*channelLink).htlcUpdates
+			select {
+			case <-secondBobChannelLink.(*channelLink).htlcUpdates:
+			case <-secondBobChannelLink.(*channelLink).quit:
+				return
+			}
 		}
 	}()
 
@@ -860,6 +890,7 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 		Cancel: func() {
 		},
 	}
+	carolTicker := time.NewTicker(50 * time.Millisecond)
 	carolChannelLink := NewChannelLink(
 		ChannelLinkConfig{
 			FwrdingPolicy:     globalPolicy,
@@ -880,6 +911,8 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 			},
 			ChainEvents: &contractcourt.ChainEventSubscription{},
 			SyncStates:  true,
+			BatchTicker: &mockTicker{carolTicker.C},
+			BatchSize:   10,
 		},
 		carolChannel,
 		startingHeight,
@@ -889,7 +922,11 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 	}
 	go func() {
 		for {
-			<-carolChannelLink.(*channelLink).htlcUpdates
+			select {
+			case <-carolChannelLink.(*channelLink).htlcUpdates:
+			case <-carolChannelLink.(*channelLink).quit:
+				return
+			}
 		}
 	}()
 
@@ -897,17 +934,21 @@ func newThreeHopNetwork(t testing.TB, aliceChannel, firstBobChannel,
 		aliceServer:      aliceServer,
 		aliceChannelLink: aliceChannelLink.(*channelLink),
 		aliceBlockEpoch:  aliceEpochChan,
+		aliceTicker:      aliceTicker,
 
 		firstBobChannelLink: firstBobChannelLink.(*channelLink),
 		bobFirstBlockEpoch:  bobFirstEpochChan,
+		firstBobTicker:      firstBobTicker,
 
 		bobServer:            bobServer,
 		secondBobChannelLink: secondBobChannelLink.(*channelLink),
 		bobSecondBlockEpoch:  bobSecondEpochChan,
+		secondBobTicker:      secondBobTicker,
 
 		carolChannelLink: carolChannelLink.(*channelLink),
 		carolServer:      carolServer,
 		carolBlockEpoch:  carolBlockEpoch,
+		carolTicker:      carolTicker,
 
 		feeEstimator: feeEstimator,
 		globalPolicy: globalPolicy,