From b4582db447a79af70d7ef36571b7e122a5cb59bf Mon Sep 17 00:00:00 2001
From: Lukasz Zimnoch <lukasz.zimnoch@keep.network>
Date: Wed, 29 Nov 2023 10:53:50 +0100
Subject: [PATCH] Fix flaky `TestNode_RunCoordinationLayer` test

---
 pkg/tbtc/node_test.go | 37 +++++++++++++++++++++++--------------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/pkg/tbtc/node_test.go b/pkg/tbtc/node_test.go
index 89649c286c..45797efc0b 100644
--- a/pkg/tbtc/node_test.go
+++ b/pkg/tbtc/node_test.go
@@ -324,13 +324,13 @@ func TestNode_RunCoordinationLayer(t *testing.T) {
 		return nil, false
 	}
 
-	// Simply add processed results to the list.
-	var processedResults []*coordinationResult
+	// Simply pass processed results to the channel.
+	processedResultsChan := make(chan *coordinationResult, 5)
 	processCoordinationResultFn := func(
 		_ *node,
 		result *coordinationResult,
 	) {
-		processedResults = append(processedResults, result)
+		processedResultsChan <- result
 	}
 
 	ctx, cancelCtx := context.WithCancel(context.Background())
@@ -347,21 +347,30 @@ func TestNode_RunCoordinationLayer(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	// Wait until the second-last coordination window passes.
-	err = localChain.blockCounter.WaitForBlockHeight(4000)
+	// Set up a stop signal that will be triggered after the last coordination
+	// window passes.
+	waiter, err := localChain.blockCounter.BlockHeightWaiter(5000)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	// Stop coordination layer. As we are between the second-last and the last
-	// coordination window, the last window should not be processed. This
-	// allows us to test that the coordination layer's shutdown works as expected.
-	cancelCtx()
-
-	// Wait until the last coordination window passes.
-	err = localChain.blockCounter.WaitForBlockHeight(5000)
-	if err != nil {
-		t.Fatal(err)
+	var processedResults []*coordinationResult
+loop:
+	for {
+		select {
+		case result := <-processedResultsChan:
+			processedResults = append(processedResults, result)
+
+			// Once the second-last coordination window is processed, stop the
+			// coordination layer. In that case, the last window should not be
+			// processed. This allows us to test that the coordination layer's
+			// shutdown works as expected.
+			if len(processedResults) == 3 {
+				cancelCtx()
+			}
+		case <-waiter:
+			break loop
+		}
 	}
 
 	testutils.AssertIntsEqual(