From 68f7642564864db0bacece93a8aa9eb563f2fb5a Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Thu, 11 Jul 2019 13:14:38 +0200
Subject: [PATCH] autopilot/agent_test: wait for expected num chans and balance

Previously we waited only for the number of channels to become what we
expected, but this wasn't enough. Sometimes the agent had't yet updated
its internal balance, causing the test to fail with an unexpected
balance.
---
 autopilot/agent_test.go | 40 +++++++++++++++++++++++-----------------
 1 file changed, 23 insertions(+), 17 deletions(-)

diff --git a/autopilot/agent_test.go b/autopilot/agent_test.go
index 92b2b358..dff8f565 100644
--- a/autopilot/agent_test.go
+++ b/autopilot/agent_test.go
@@ -1259,34 +1259,40 @@ func TestAgentChannelSizeAllocation(t *testing.T) {
 	waitForNumChans := func(expChans int) {
 		t.Helper()
 
+		var (
+			numChans int
+			balance  btcutil.Amount
+		)
+
 	Loop:
 		for {
 			select {
 			case arg := <-testCtx.constraints.moreChanArgs:
+				numChans = len(arg.chans)
+				balance = arg.balance
+
 				// As long as the number of existing channels
 				// is below our expected number of channels,
-				// we'll keep responding with "no more
-				// channels".
-				if len(arg.chans) != expChans {
-					select {
-					case testCtx.constraints.moreChansResps <- moreChansResp{0, 0}:
-					case <-time.After(time.Second * 3):
-						t.Fatalf("heuristic wasn't " +
-							"queried in time")
-					}
-					continue
+				// and the balance is not what we expect, we'll
+				// keep responding with "no more channels".
+				if numChans == expChans &&
+					balance == testCtx.walletBalance {
+					break Loop
 				}
 
-				if arg.balance != testCtx.walletBalance {
-					t.Fatalf("expectd agent to have %v "+
-						"balance, had %v",
-						testCtx.walletBalance,
-						arg.balance)
+				select {
+				case testCtx.constraints.moreChansResps <- moreChansResp{0, 0}:
+				case <-time.After(time.Second * 3):
+					t.Fatalf("heuristic wasn't queried " +
+						"in time")
 				}
-				break Loop
 
 			case <-time.After(time.Second * 3):
-				t.Fatalf("heuristic wasn't queried in time")
+				t.Fatalf("did not receive expected "+
+					"channels(%d) and balance(%d), "+
+					"instead got %d and %d", expChans,
+					testCtx.walletBalance, numChans,
+					balance)
 			}
 		}
 	}