From 7838e6e1fa4d2483cfa1c3a4dc41da44032189ae Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Thu, 11 Jul 2019 13:14:39 +0200
Subject: [PATCH] autopilot/agent_test: add more agent channel allocation
 testcases

Add tests for allocation when well within and well above channel budget.
---
 autopilot/agent_test.go | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/autopilot/agent_test.go b/autopilot/agent_test.go
index dff8f565..bd703c9e 100644
--- a/autopilot/agent_test.go
+++ b/autopilot/agent_test.go
@@ -1196,7 +1196,7 @@ func TestAgentChannelSizeAllocation(t *testing.T) {
 	t.Parallel()
 
 	// Total number of nodes in our mock graph.
-	const numNodes = 10
+	const numNodes = 20
 
 	testCtx, cleanup := setup(t, nil)
 	defer cleanup()
@@ -1240,6 +1240,8 @@ func TestAgentChannelSizeAllocation(t *testing.T) {
 		numNewChannels, nodeScores,
 	)
 
+	// We expect the autopilot to have allocated all funds towards
+	// channels.
 	expectedAllocation := testCtx.constraints.MaxChanSize() * btcutil.Amount(numNewChannels)
 	nodes := checkChannelOpens(
 		t, testCtx, expectedAllocation, numNewChannels,
@@ -1314,5 +1316,30 @@ func TestAgentChannelSizeAllocation(t *testing.T) {
 
 	// To stay within the budget, we expect the autopilot to open 2
 	// channels.
-	checkChannelOpens(t, testCtx, channelBudget, 2)
+	expectedAllocation = channelBudget
+	nodes = checkChannelOpens(t, testCtx, expectedAllocation, 2)
+	numExistingChannels = 7
+
+	for _, node := range nodes {
+		delete(nodeScores, node)
+	}
+
+	waitForNumChans(numExistingChannels)
+
+	// Finally check that we make maximum channels if we are well within
+	// our budget.
+	channelBudget = btcutil.SatoshiPerBitcoin * 5
+	numNewChannels = 2
+	respondWithScores(
+		t, testCtx, channelBudget, numExistingChannels,
+		numNewChannels, nodeScores,
+	)
+
+	// We now expect the autopilot to open 2 channels, and since it has
+	// more than enough balance within the budget, they should both be of
+	// maximum size.
+	expectedAllocation = testCtx.constraints.MaxChanSize() *
+		btcutil.Amount(numNewChannels)
+
+	checkChannelOpens(t, testCtx, expectedAllocation, numNewChannels)
 }