Ensure calls to Node.Stop() provide context

ava-labs · Nov 27, 2023 · 791b12d · 791b12d
1 parent afad840
commit 791b12d
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 9 deletions.
diff --git a/tests/fixture/e2e/helpers.go b/tests/fixture/e2e/helpers.go
@@ -143,7 +143,7 @@ func DefaultEventually(condition func() bool, msg string) {
 func AddEphemeralNode(network testnet.Network, flags testnet.FlagsMap) testnet.Node {
 	require := require.New(ginkgo.GinkgoT())
 
-	node, err := network.AddEphemeralNode(ginkgo.GinkgoWriter, flags)
+	node, err := network.AddEphemeralNode(DefaultContext(), ginkgo.GinkgoWriter, flags)
 	require.NoError(err)
 
 	RegisterNodeforCleanup(node)
@@ -217,12 +217,16 @@ func CheckBootstrapIsPossible(network testnet.Network) {
 	// checking for bootstrap implicitly on teardown via a function registered
 	// with ginkgo.DeferCleanup. It's not possible to call DeferCleanup from
 	// within a function called by DeferCleanup.
-	node, err := network.AddEphemeralNode(ginkgo.GinkgoWriter, testnet.FlagsMap{})
+	ctx, cancel := context.WithTimeout(context.Background(), DefaultTimeout)
+	defer cancel()
+	node, err := network.AddEphemeralNode(ctx, ginkgo.GinkgoWriter, testnet.FlagsMap{})
 	require.NoError(err)
 
 	defer func() {
 		tests.Outf("Shutting down ephemeral node %s\n", node.GetID())
-		require.NoError(node.Stop())
+		ctx, cancel := context.WithTimeout(context.Background(), DefaultTimeout)
+		defer cancel()
+		require.NoError(node.Stop(ctx, true /* waitForProcessStopped */))
 	}()
 
 	WaitForHealthy(node)

diff --git a/tests/fixture/testnet/interfaces.go b/tests/fixture/testnet/interfaces.go
@@ -15,7 +15,7 @@ import (
 type Network interface {
 	GetConfig() NetworkConfig
 	GetNodes() []Node
-	AddEphemeralNode(w io.Writer, flags FlagsMap) (Node, error)
+	AddEphemeralNode(ctx context.Context, w io.Writer, flags FlagsMap) (Node, error)
 	GetEphemeralNodes(nodeIDs []ids.NodeID) ([]Node, error)
 	GetSubnets() ([]*Subnet, error)
 	WriteSubnets([]*Subnet) error

diff --git a/tests/fixture/testnet/local/network.go b/tests/fixture/testnet/local/network.go
@@ -126,14 +126,14 @@ func (ln *LocalNetwork) GetNodes() []testnet.Node {
 }
 
 // Adds a backend-agnostic ephemeral node to the network
-func (ln *LocalNetwork) AddEphemeralNode(w io.Writer, flags testnet.FlagsMap) (testnet.Node, error) {
+func (ln *LocalNetwork) AddEphemeralNode(ctx context.Context, w io.Writer, flags testnet.FlagsMap) (testnet.Node, error) {
 	if flags == nil {
 		flags = testnet.FlagsMap{}
 	} else {
 		// Avoid modifying the input flags map
 		flags = flags.Copy()
 	}
-	return ln.AddLocalNode(w, &LocalNode{
+	return ln.AddLocalNode(ctx, w, &LocalNode{
 		NodeConfig: testnet.NodeConfig{
 			Flags: flags,
 		},
@@ -689,7 +689,7 @@ func (ln *LocalNetwork) ReadAll() error {
 	return ln.ReadNodes()
 }
 
-func (ln *LocalNetwork) AddLocalNode(w io.Writer, node *LocalNode, isEphemeral bool) (*LocalNode, error) {
+func (ln *LocalNetwork) AddLocalNode(ctx context.Context, w io.Writer, node *LocalNode, isEphemeral bool) (*LocalNode, error) {
 	// Assume network configuration has been written to disk and is current in memory
 
 	if node == nil {
@@ -736,7 +736,7 @@ func (ln *LocalNetwork) AddLocalNode(w io.Writer, node *LocalNode, isEphemeral b
 	if err != nil {
 		// Attempt to stop an unhealthy node to provide some assurance to the caller
 		// that an error condition will not result in a lingering process.
-		stopErr := node.Stop()
+		stopErr := node.Stop(ctx, true /* waitForProcessStopped */)
 		if stopErr != nil {
 			err = errors.Join(err, stopErr)
 		}

diff --git a/tests/fixture/testnet/subnets.go b/tests/fixture/testnet/subnets.go
@@ -256,7 +256,7 @@ func CreateSubnet(
 		flags := nodeSpec.Flags.Copy()
 		flags[config.TrackSubnetsKey] = subnetID.String()
 		for i := 0; i < nodeSpec.Count; i++ {
-			node, err := network.AddEphemeralNode(w, flags)
+			node, err := network.AddEphemeralNode(ctx, w, flags)
 			if err != nil {
 				return nil, err
 			}