From 347d1545de01dad50afc770270c3c3336969ed34 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 19 Dec 2018 14:14:53 +0100
Subject: [PATCH] utxonursery: use mocked sweeper in tests

This commit fixes a test flake caused by a race condition. Using the
real sweeper in the nursery test created too complex concurrent
behaviour to reliably assert on and made the tests difficult to
comprehend.
---
 sweep/test_utils.go |  6 ++--
 utxonursery_test.go | 86 ++++++---------------------------------------
 2 files changed, 13 insertions(+), 79 deletions(-)

diff --git a/sweep/test_utils.go b/sweep/test_utils.go
index b704217f..d9ef157c 100644
--- a/sweep/test_utils.go
+++ b/sweep/test_utils.go
@@ -2,8 +2,6 @@ package sweep
 
 import (
 	"fmt"
-	"os"
-	"runtime/pprof"
 	"sync"
 	"testing"
 	"time"
@@ -58,6 +56,8 @@ func NewMockNotifier(t *testing.T) *MockNotifier {
 
 // NotifyEpoch simulates a new epoch arriving.
 func (m *MockNotifier) NotifyEpoch(height int32) {
+	m.t.Helper()
+
 	for epochChan, chanHeight := range m.epochChan {
 		// Only send notifications if the height is greater than the
 		// height the caller passed into the register call.
@@ -72,8 +72,6 @@ func (m *MockNotifier) NotifyEpoch(height int32) {
 			Height: height,
 		}:
 		case <-time.After(defaultTestTimeout):
-			pprof.Lookup("goroutine").WriteTo(os.Stdout, 1)
-
 			m.t.Fatal("epoch event not consumed")
 		}
 	}
diff --git a/utxonursery_test.go b/utxonursery_test.go
index a6eef7be..2f8629b5 100644
--- a/utxonursery_test.go
+++ b/utxonursery_test.go
@@ -403,7 +403,7 @@ type nurseryTestContext struct {
 	store       *nurseryStoreInterceptor
 	restart     func() bool
 	receiveTx   func() wire.MsgTx
-	sweeper     *sweep.UtxoSweeper
+	sweeper     *mockSweeper
 	timeoutChan chan chan time.Time
 	t           *testing.T
 }
@@ -449,33 +449,7 @@ func createNurseryTestContext(t *testing.T,
 		bestHeight: 0,
 	}
 
-	sweeperStore := sweep.NewMockSweeperStore()
-
-	sweeperCfg := &sweep.UtxoSweeperConfig{
-		GenSweepScript: func() ([]byte, error) {
-			return []byte{}, nil
-		},
-		Estimator: &lnwallet.StaticFeeEstimator{},
-		Signer:    &nurseryMockSigner{},
-		Notifier:  notifier,
-		PublishTransaction: func(tx *wire.MsgTx) error {
-			return publishFunc(tx, "sweeper")
-		},
-		NewBatchTimer: func() <-chan time.Time {
-			c := make(chan time.Time, 1)
-			timeoutChan <- c
-			return c
-		},
-		ChainIO:              chainIO,
-		Store:                sweeperStore,
-		MaxInputsPerTx:       10,
-		MaxSweepAttempts:     5,
-		NextAttemptDeltaFunc: func(int) int32 { return 1 },
-	}
-
-	sweeper := sweep.New(sweeperCfg)
-
-	sweeper.Start()
+	sweeper := newMockSweeper(t)
 
 	nurseryCfg := NurseryConfig{
 		Notifier: notifier,
@@ -491,7 +465,7 @@ func createNurseryTestContext(t *testing.T,
 		},
 		Store:      storeIntercepter,
 		ChainIO:    chainIO,
-		SweepInput: sweeper.SweepInput,
+		SweepInput: sweeper.sweepInput,
 		PublishTransaction: func(tx *wire.MsgTx) error {
 			return publishFunc(tx, "nursery")
 		},
@@ -531,33 +505,11 @@ func createNurseryTestContext(t *testing.T,
 			// Simulate lnd restart.
 			ctx.nursery.Stop()
 
-			// Also restart sweeper to test behaviour as one unit.
-			//
-			// TODO(joostjager): Mock sweeper to test nursery in
-			// isolation.
-			ctx.sweeper.Stop()
-
-			// Find out if there is a last tx stored. If so, we
-			// expect it to be republished on startup.
-			hasLastTx, err := sweeperCfg.Store.GetLastPublishedTx()
-			if err != nil {
-				t.Fatal(err)
-			}
-
 			// Restart sweeper.
-			ctx.sweeper = sweep.New(sweeperCfg)
-			ctx.sweeper.Start()
-
-			// Receive last tx if expected.
-			if hasLastTx != nil {
-				utxnLog.Debugf("Expecting republish")
-				ctx.receiveTx()
-			} else {
-				utxnLog.Debugf("Expecting no republish")
-			}
+			ctx.sweeper = newMockSweeper(t)
 
 			/// Restart nursery.
-			nurseryCfg.SweepInput = ctx.sweeper.SweepInput
+			nurseryCfg.SweepInput = ctx.sweeper.sweepInput
 			ctx.nursery = newUtxoNursery(&nurseryCfg)
 			ctx.nursery.Start()
 
@@ -571,6 +523,8 @@ func createNurseryTestContext(t *testing.T,
 }
 
 func (ctx *nurseryTestContext) notifyEpoch(height int32) {
+	ctx.t.Helper()
+
 	ctx.chainIO.bestHeight = height
 	ctx.notifier.NotifyEpoch(height)
 }
@@ -630,8 +584,6 @@ func (ctx *nurseryTestContext) finish() {
 	if len(activeHeights) > 0 {
 		ctx.t.Fatalf("Expected height index to be empty")
 	}
-
-	ctx.sweeper.Stop()
 }
 
 func createOutgoingRes(onLocalCommitment bool) *lnwallet.OutgoingHtlcResolution {
@@ -956,20 +908,17 @@ func testSweep(t *testing.T, ctx *nurseryTestContext,
 	afterPublishAssert func()) {
 
 	// Wait for nursery to publish the sweep tx.
-	ctx.tick()
-	sweepTx := ctx.receiveTx()
+	ctx.sweeper.expectSweep()
 
 	if ctx.restart() {
-		// Nursery reoffers its input. Sweeper needs a tick to create the sweep
-		// tx.
-		ctx.tick()
-		ctx.receiveTx()
+		// Nursery reoffers its input after a restart.
+		ctx.sweeper.expectSweep()
 	}
 
 	afterPublishAssert()
 
 	// Confirm the sweep tx.
-	ctx.notifier.SpendOutpoint(sweepTx.TxIn[0].PreviousOutPoint, sweepTx)
+	ctx.sweeper.sweepAll()
 
 	// Wait for output to be promoted in store to GRAD.
 	select {
@@ -986,19 +935,6 @@ func testSweep(t *testing.T, ctx *nurseryTestContext,
 	assertNurseryReportUnavailable(t, ctx.nursery)
 }
 
-func (ctx *nurseryTestContext) tick() {
-	select {
-	case c := <-ctx.timeoutChan:
-		select {
-		case c <- time.Time{}:
-		case <-time.After(defaultTestTimeout):
-			ctx.t.Fatal("tick timeout - tick not consumed")
-		}
-	case <-time.After(defaultTestTimeout):
-		ctx.t.Fatal("tick timeout - no new timer created")
-	}
-}
-
 type nurseryStoreInterceptor struct {
 	ns NurseryStore