From cf3dd3fb945b57505bae7c50c50538c02746127e Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Mon, 3 Jun 2019 20:32:07 +0200
Subject: [PATCH] routing/test: create mission control test context

---
 routing/missioncontrol_test.go | 87 ++++++++++++++++++++++------------
 1 file changed, 57 insertions(+), 30 deletions(-)

diff --git a/routing/missioncontrol_test.go b/routing/missioncontrol_test.go
index 19a28828..72c3b07a 100644
--- a/routing/missioncontrol_test.go
+++ b/routing/missioncontrol_test.go
@@ -8,9 +8,25 @@ import (
 	"github.com/lightningnetwork/lnd/routing/route"
 )
 
-// TestMissionControl tests mission control probability estimation.
-func TestMissionControl(t *testing.T) {
-	now := testTime
+var (
+	mcTestNode = route.Vertex{}
+	mcTestEdge = EdgeLocator{
+		ChannelID: 123,
+	}
+	mcTestTime = time.Date(2018, time.January, 9, 14, 00, 00, 0, time.UTC)
+)
+
+type mcTestContext struct {
+	t   *testing.T
+	mc  *MissionControl
+	now time.Time
+}
+
+func createMcTestContext(t *testing.T) *mcTestContext {
+	ctx := &mcTestContext{
+		t:   t,
+		now: mcTestTime,
+	}
 
 	mc := NewMissionControl(
 		nil, nil, nil, &MissionControlConfig{
@@ -18,7 +34,30 @@ func TestMissionControl(t *testing.T) {
 			AprioriHopProbability: 0.8,
 		},
 	)
-	mc.now = func() time.Time { return now }
+
+	mc.now = func() time.Time { return ctx.now }
+	ctx.mc = mc
+
+	return ctx
+}
+
+// Assert that mission control returns a probability for an edge.
+func (ctx *mcTestContext) expectP(amt lnwire.MilliSatoshi,
+	expected float64) {
+
+	ctx.t.Helper()
+
+	p := ctx.mc.getEdgeProbability(mcTestNode, mcTestEdge, amt)
+	if p != expected {
+		ctx.t.Fatalf("unexpected probability %v", p)
+	}
+}
+
+// TestMissionControl tests mission control probability estimation.
+func TestMissionControl(t *testing.T) {
+	ctx := createMcTestContext(t)
+
+	ctx.now = testTime
 
 	testTime := time.Date(2018, time.January, 9, 14, 00, 00, 0, time.UTC)
 
@@ -27,50 +66,38 @@ func TestMissionControl(t *testing.T) {
 		channel: 123,
 	}
 
-	expectP := func(amt lnwire.MilliSatoshi, expected float64) {
-		t.Helper()
-
-		p := mc.getEdgeProbability(
-			testNode, EdgeLocator{ChannelID: testEdge.channel},
-			amt,
-		)
-		if p != expected {
-			t.Fatalf("unexpected probability %v", p)
-		}
-	}
-
 	// Initial probability is expected to be 1.
-	expectP(1000, 0.8)
+	ctx.expectP(1000, 0.8)
 
 	// Expect probability to be zero after reporting the edge as failed.
-	mc.reportEdgeFailure(testEdge, 1000)
-	expectP(1000, 0)
+	ctx.mc.reportEdgeFailure(testEdge, 1000)
+	ctx.expectP(1000, 0)
 
 	// As we reported with a min penalization amt, a lower amt than reported
 	// should be unaffected.
-	expectP(500, 0.8)
+	ctx.expectP(500, 0.8)
 
 	// Edge decay started.
-	now = testTime.Add(30 * time.Minute)
-	expectP(1000, 0.4)
+	ctx.now = testTime.Add(30 * time.Minute)
+	ctx.expectP(1000, 0.4)
 
 	// Edge fails again, this time without a min penalization amt. The edge
 	// should be penalized regardless of amount.
-	mc.reportEdgeFailure(testEdge, 0)
-	expectP(1000, 0)
-	expectP(500, 0)
+	ctx.mc.reportEdgeFailure(testEdge, 0)
+	ctx.expectP(1000, 0)
+	ctx.expectP(500, 0)
 
 	// Edge decay started.
-	now = testTime.Add(60 * time.Minute)
-	expectP(1000, 0.4)
+	ctx.now = testTime.Add(60 * time.Minute)
+	ctx.expectP(1000, 0.4)
 
 	// A node level failure should bring probability of every channel back
 	// to zero.
-	mc.reportVertexFailure(testNode)
-	expectP(1000, 0)
+	ctx.mc.reportVertexFailure(testNode)
+	ctx.expectP(1000, 0)
 
 	// Check whether history snapshot looks sane.
-	history := mc.GetHistorySnapshot()
+	history := ctx.mc.GetHistorySnapshot()
 	if len(history.Nodes) != 1 {
 		t.Fatal("unexpected number of nodes")
 	}