From 2731409f619ade8a00162cb91ebbad38be7dc463 Mon Sep 17 00:00:00 2001
From: Maru Newby <maru.newby@avalabs.org>
Date: Thu, 4 Jan 2024 14:36:07 -0800
Subject: [PATCH] tmpnet: Retain IDs of validating nodes with subnets

---
 tests/fixture/e2e/env.go        | 20 +++++++++++++-------
 tests/fixture/tmpnet/network.go |  2 +-
 tests/fixture/tmpnet/subnet.go  | 10 +++++++---
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/tests/fixture/e2e/env.go b/tests/fixture/e2e/env.go
index 3e6c67dbe4b8..567c8f1f05bf 100644
--- a/tests/fixture/e2e/env.go
+++ b/tests/fixture/e2e/env.go
@@ -86,15 +86,21 @@ func NewTestEnvironment(flagVars *FlagVars, desiredNetwork *tmpnet.Network) *Tes
 	require.NoError(network.CreateSubnets(DefaultContext(), ginkgo.GinkgoWriter))
 
 	// Wait for chains to have bootstrapped on all nodes
+	nodeURIs := network.GetNodeURIs()
 	Eventually(func() bool {
 		for _, subnet := range network.Subnets {
-			for _, node := range network.Nodes {
-				infoClient := info.NewClient(node.URI)
-				for _, chain := range subnet.Chains {
-					isBootstrapped, err := infoClient.IsBootstrapped(DefaultContext(), chain.ChainID.String())
-					// Ignore errors since a chain id that is not yet known will result in a recoverable error.
-					if err != nil || !isBootstrapped {
-						return false
+			for _, validatorID := range subnet.ValidatorIDs {
+				for _, nodeURI := range nodeURIs {
+					if validatorID == nodeURI.NodeID {
+						infoClient := info.NewClient(nodeURI.URI)
+						for _, chain := range subnet.Chains {
+							isBootstrapped, err := infoClient.IsBootstrapped(DefaultContext(), chain.ChainID.String())
+							// Ignore errors since a chain id that is not yet known will result in a recoverable error.
+							if err != nil || !isBootstrapped {
+								return false
+							}
+						}
+						break
 					}
 				}
 			}
diff --git a/tests/fixture/tmpnet/network.go b/tests/fixture/tmpnet/network.go
index f39c25c07eb7..ee1f0d8b8c84 100644
--- a/tests/fixture/tmpnet/network.go
+++ b/tests/fixture/tmpnet/network.go
@@ -577,7 +577,7 @@ func (n *Network) CreateSubnets(ctx context.Context, w io.Writer) error {
 	// Wait for nodes to become subnet validators
 	pChainClient := platformvm.NewClient(n.Nodes[0].URI)
 	for _, subnet := range createdSubnets {
-		if err := waitForActiveValidators(ctx, w, pChainClient, subnet, n.Nodes); err != nil {
+		if err := waitForActiveValidators(ctx, w, pChainClient, subnet); err != nil {
 			return err
 		}
 
diff --git a/tests/fixture/tmpnet/subnet.go b/tests/fixture/tmpnet/subnet.go
index da253499d825..89e55268c61c 100644
--- a/tests/fixture/tmpnet/subnet.go
+++ b/tests/fixture/tmpnet/subnet.go
@@ -68,6 +68,9 @@ type Subnet struct {
 	// The private key that owns the subnet
 	OwningKey *secp256k1.PrivateKey
 
+	// IDs of the nodes responsible for validating the subnet
+	ValidatorIDs []ids.NodeID
+
 	Chains []*Chain
 }
 
@@ -188,6 +191,8 @@ func (s *Subnet) AddValidators(ctx context.Context, nodes []*Node) error {
 		if err != nil {
 			return err
 		}
+
+		s.ValidatorIDs = append(s.ValidatorIDs, node.NodeID)
 	}
 
 	return nil
@@ -236,7 +241,6 @@ func waitForActiveValidators(
 	w io.Writer,
 	pChainClient platformvm.Client,
 	subnet *Subnet,
-	nodes []*Node,
 ) error {
 	ticker := time.NewTicker(DefaultPollingInterval)
 	defer ticker.Stop()
@@ -258,8 +262,8 @@ func waitForActiveValidators(
 			validatorSet.Add(validator.NodeID)
 		}
 		allActive := true
-		for _, node := range nodes {
-			if !validatorSet.Contains(node.NodeID) {
+		for _, validatorID := range subnet.ValidatorIDs {
+			if !validatorSet.Contains(validatorID) {
 				allActive = false
 			}
 		}