Extract common channel handling to helper func

ava-labs · Aug 16, 2024 · c182fa6 · c182fa6
1 parent fa3d590
commit c182fa6
Show file tree

Hide file tree

Showing 9 changed files with 27 additions and 80 deletions.
diff --git a/tests/allowed_addresses.go b/tests/allowed_addresses.go
@@ -177,12 +177,7 @@ func AllowedAddresses(network interfaces.LocalNetwork) {
 
 	// Wait for relayer to start up
 	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-		close(readyChan)
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	// Allowed by Relayer 1
 	testUtils.RelayBasicMessage(
@@ -206,12 +201,7 @@ func AllowedAddresses(network interfaces.LocalNetwork) {
 
 	// Wait for relayer to start up
 	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-		close(readyChan)
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	// Disallowed by Relayer 2
 	_, _, id := testUtils.SendBasicTeleporterMessage(
@@ -251,12 +241,7 @@ func AllowedAddresses(network interfaces.LocalNetwork) {
 
 	// Wait for relayer to start up
 	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-		close(readyChan)
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	// Disallowed by Relayer 3
 	_, _, id = testUtils.SendBasicTeleporterMessage(
@@ -295,13 +280,7 @@ func AllowedAddresses(network interfaces.LocalNetwork) {
 	defer relayerCleanup()
 
 	// Wait for relayer to start up
-	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-		close(readyChan)
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	// Disallowed by Relayer 4
 	_, _, id = testUtils.SendBasicTeleporterMessage(

diff --git a/tests/basic_relay.go b/tests/basic_relay.go
@@ -67,12 +67,7 @@ func BasicRelay(network interfaces.LocalNetwork) {
 	defer relayerCleanup()
 
 	// Wait for relayer to start up
-	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	log.Info("Sending transaction from Subnet A to Subnet B")
 	testUtils.RelayBasicMessage(
@@ -152,12 +147,7 @@ func BasicRelay(network interfaces.LocalNetwork) {
 
 	// Wait for relayer to start up
 	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-		close(readyChan)
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	// We should not receive a new block on subnet B, since the relayer should have
 	// seen the Teleporter message was already delivered.

diff --git a/tests/batch_relay.go b/tests/batch_relay.go
@@ -72,12 +72,7 @@ func BatchRelay(network interfaces.LocalNetwork) {
 
 	// Wait for relayer to start up
 	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-		close(readyChan)
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	//
 	// Send a batch message from subnet A -> B

diff --git a/tests/manual_message.go b/tests/manual_message.go
@@ -114,11 +114,7 @@ func ManualMessage(network interfaces.LocalNetwork) {
 
 	// Wait for relayer to startup.
 	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	reqBody := api.ManualWarpMessageRequest{
 		UnsignedMessageBytes: unsignedMessage.Bytes(),

diff --git a/tests/relay_message_api.go b/tests/relay_message_api.go
@@ -71,12 +71,7 @@ func RelayMessageAPI(network interfaces.LocalNetwork) {
 
 	// Wait for relayer to start up
 	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-		close(readyChan)
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	reqBody := api.RelayMessageRequest{
 		BlockchainID: subnetAInfo.BlockchainID.String(),

diff --git a/tests/shared_db.go b/tests/shared_db.go
@@ -77,8 +77,12 @@ func SharedDatabaseAccess(network interfaces.LocalNetwork) {
 	log.Info("Waiting for the relayers to start up")
 	var wg sync.WaitGroup
 	wg.Add(2)
-	go testUtils.WaitFunc(&wg, readyChanA)
-	go testUtils.WaitFunc(&wg, readyChanB)
+	waitFunc := func(wg *sync.WaitGroup, readyChan chan struct{}) {
+		defer wg.Done()
+		<-readyChan
+	}
+	go waitFunc(&wg, readyChanA)
+	go waitFunc(&wg, readyChanB)
 	wg.Wait()
 
 	log.Info("Sending transaction from Subnet A to Subnet B")

diff --git a/tests/signature_aggregator_api.go b/tests/signature_aggregator_api.go
@@ -57,12 +57,7 @@ func SignatureAggregatorAPI(network interfaces.LocalNetwork) {
 
 	// Wait for signature-aggregator to start up
 	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-		close(readyChan)
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Signature Aggregator did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	// End setup step
 	// Begin Test Case 1

diff --git a/tests/utils/utils.go b/tests/utils/utils.go
@@ -14,7 +14,6 @@ import (
 	"os"
 	"os/exec"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/ava-labs/avalanchego/ids"
@@ -490,13 +489,7 @@ func TriggerProcessMissedBlocks(
 	defer relayerCleanup()
 
 	// Wait for relayer to start up
-	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-		close(readyChan)
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	WaitForChannelClose(readyChan, 15*time.Second)
 
 	log.Info("Waiting for a new block confirmation on the destination")
 	<-newHeads
@@ -566,7 +559,7 @@ func runExecutable(
 		for scanner.Scan() {
 			text := scanner.Text()
 			if strings.Contains(text, "Initialization complete") {
-				readyChan <- struct{}{}
+				close(readyChan)
 			}
 			log.Info(text)
 		}
@@ -590,7 +583,12 @@ func runExecutable(
 	return readyChan
 }
 
-func WaitFunc(wg *sync.WaitGroup, ch chan struct{}) {
-	defer wg.Done()
-	<-ch
+// Helper function that waits for a signaling channel to be closed
+// or throws an error if the channel is not closed in time
+func WaitForChannelClose(ch <-chan struct{}, timeout time.Duration) {
+	select {
+	case <-ch:
+	case <-time.After(timeout):
+		Expect(false).To(BeTrue(), "Channel did not close in time")
+	}
 }
diff --git a/tests/warp_api.go b/tests/warp_api.go
@@ -74,12 +74,7 @@ func WarpAPIRelay(network interfaces.LocalNetwork) {
 
 	// Wait for relayer to start up
 	log.Info("Waiting for the relayer to start up")
-	select {
-	case <-readyChan:
-		close(readyChan)
-	case <-time.After(15 * time.Second):
-		Expect(false).To(BeTrue(), "Relayer did not start up in time")
-	}
+	testUtils.WaitForChannelClose(readyChan, 15*time.Second)
 
 	log.Info("Sending transaction from Subnet A to Subnet B")
 	testUtils.RelayBasicMessage(