Pass external context to NewLocalNetwork and remove hardcoded timeout…

…s inside the package
ava-labs · Sep 6, 2024 · 97dae43 · 97dae43
1 parent d15b1f8
commit 97dae43
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 10 deletions.
diff --git a/tests/flows/teleporter_registry.go b/tests/flows/teleporter_registry.go
@@ -2,6 +2,7 @@ package flows
 
 import (
 	"context"
+	"time"
 
 	"github.com/ava-labs/subnet-evm/accounts/abi/bind"
 	"github.com/ava-labs/teleporter/tests/interfaces"
@@ -76,7 +77,9 @@ func TeleporterRegistry(network interfaces.LocalNetwork) {
 
 	// Restart nodes with new chain config
 	network.SetChainConfigs(chainConfigs)
-	network.RestartNodes(ctx, network.GetAllNodeIDs())
+	restartCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+	network.RestartNodes(restartCtx, network.GetAllNodeIDs())
 
 	// Call addProtocolVersion on subnetB to register the new Teleporter version
 	utils.AddProtocolVersionAndWaitForAcceptance(

diff --git a/tests/flows/validator_churn.go b/tests/flows/validator_churn.go
@@ -3,6 +3,7 @@ package flows
 import (
 	"context"
 	"math/big"
+	"time"
 
 	"github.com/ava-labs/subnet-evm/accounts/abi/bind"
 	subnetEvmUtils "github.com/ava-labs/subnet-evm/tests/utils"
@@ -59,7 +60,9 @@ func ValidatorChurn(network interfaces.LocalNetwork) {
 	//
 
 	// Add new nodes to the validator set
-	network.AddSubnetValidators(ctx, subnetAInfo.SubnetID, newNodeCount)
+	addValidatorsCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
+	defer cancel()
+	network.AddSubnetValidators(addValidatorsCtx, subnetAInfo.SubnetID, newNodeCount)
 
 	// Refresh the subnet info
 	subnetAInfo, subnetBInfo = utils.GetTwoSubnets(network)

diff --git a/tests/flows/validator_set_sig.go b/tests/flows/validator_set_sig.go
@@ -3,6 +3,7 @@ package flows
 import (
 	"context"
 	"math/big"
+	"time"
 
 	"github.com/ava-labs/subnet-evm/accounts/abi/bind"
 	validatorsetsig "github.com/ava-labs/teleporter/abi-bindings/go/governance/ValidatorSetSig"
@@ -135,7 +136,9 @@ func ValidatorSetSig(network interfaces.LocalNetwork) {
 
 	// Restart nodes with new chain config
 	network.SetChainConfigs(chainConfigs)
-	network.RestartNodes(ctx, network.GetAllNodeIDs())
+	restartCtx, cancel := context.WithTimeout(ctx, time.Second*30)
+	defer cancel()
+	network.RestartNodes(restartCtx, network.GetAllNodeIDs())
 
 	// ************************************************************************************************
 	// Test Case 1: validatorChain (subnetB) != targetChain (subnetA)

diff --git a/tests/local/e2e_test.go b/tests/local/e2e_test.go
@@ -4,8 +4,10 @@
 package local
 
 import (
+	"context"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/ava-labs/teleporter/tests/flows"
 	deploymentUtils "github.com/ava-labs/teleporter/utils/deployment-utils"
@@ -49,7 +51,10 @@ var _ = ginkgo.BeforeSuite(func() {
 	Expect(err).Should(BeNil())
 
 	// Create the local network instance
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
 	LocalNetworkInstance = NewLocalNetwork(
+		ctx,
 		"teleporter-test-local-network",
 		warpGenesisTemplateFile,
 		[]SubnetSpec{

diff --git a/tests/local/network.go b/tests/local/network.go
@@ -63,8 +63,6 @@ const (
 					"internal-debug","internal-account","internal-personal",
 					"debug","debug-tracer","debug-file-tracer","debug-handler"]
 	}`
-
-	timeout = 180 * time.Second
 )
 
 type SubnetSpec struct {
@@ -77,12 +75,12 @@ type SubnetSpec struct {
 }
 
 func NewLocalNetwork(
+	ctx context.Context,
 	name string,
 	warpGenesisTemplateFile string,
 	subnetSpecs []SubnetSpec,
 	extraNodeCount int, // for use by tests, eg to add new subnet validators
 ) *LocalNetwork {
-	ctx := context.Background()
 	var err error
 
 	// Create extra nodes to be used to add more validators later
@@ -128,9 +126,9 @@ func NewLocalNetwork(
 	avalancheGoBuildPath, ok := os.LookupEnv("AVALANCHEGO_BUILD_PATH")
 	Expect(ok).Should(Equal(true))
 
-	ctxWithTimeout, cancelBootstrap := context.WithTimeout(ctx, timeout)
+	ctx, cancelBootstrap := context.WithCancel(ctx)
 	err = tmpnet.BootstrapNewNetwork(
-		ctxWithTimeout,
+		ctx,
 		os.Stdout,
 		network,
 		"",
@@ -509,7 +507,7 @@ func (n *LocalNetwork) AddSubnetValidators(ctx context.Context, subnetID ids.ID,
 	apiURI, err := n.tmpnet.GetURIForNodeID(subnet.ValidatorIDs[0])
 	Expect(err).Should(BeNil())
 
-	ctx, cancel := context.WithTimeout(ctx, timeout)
+	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 	err = subnet.AddValidators(
 		ctx,
@@ -556,7 +554,7 @@ func (n *LocalNetwork) RestartNodes(ctx context.Context, nodeIDs []ids.NodeID) {
 	}
 
 	for _, node := range nodes {
-		ctx, cancel := context.WithTimeout(ctx, timeout)
+		ctx, cancel := context.WithCancel(ctx)
 		defer cancel()
 		err := node.SaveAPIPort()
 		Expect(err).Should(BeNil())