[tmpnet] Unify start network flag usage between e2e and tmpnetctl (#3871)

Author: maru-ava

tests/antithesis/config.go

@@ -89,7 +89,9 @@ func configForNewNetwork(
 	duration time.Duration,
 ) *Config {
 	if defaultNetwork.Nodes == nil {
-		defaultNetwork.Nodes = tmpnet.NewNodesOrPanic(flagVars.NodeCount())
+		nodeCount, err := flagVars.NodeCount()
+		require.NoError(tc, err)
+		defaultNetwork.Nodes = tmpnet.NewNodesOrPanic(nodeCount)
 	}
 	if defaultNetwork.Subnets == nil && getSubnets != nil {
 		defaultNetwork.Subnets = getSubnets(defaultNetwork.Nodes...)

tests/e2e/e2e_test.go

@@ -34,15 +34,17 @@ func TestE2E(t *testing.T) {
 var flagVars *e2e.FlagVars
 
 func init() {
-	flagVars = e2e.RegisterFlags()
+	flagVars = e2e.RegisterFlagsWithDefaultOwner("avalanchego-e2e")
 }
 
 var _ = ginkgo.SynchronizedBeforeSuite(func() []byte {
 	// Run only once in the first ginkgo process
 
 	tc := e2e.NewEventHandlerTestContext()
 
-	nodes := tmpnet.NewNodesOrPanic(flagVars.NodeCount())
+	nodeCount, err := flagVars.NodeCount()
+	require.NoError(tc, err)
+	nodes := tmpnet.NewNodesOrPanic(nodeCount)
 	subnets := vms.XSVMSubnetsOrPanic(nodes...)
 
 	upgrades := upgrade.Default
@@ -63,7 +65,7 @@ var _ = ginkgo.SynchronizedBeforeSuite(func() []byte {
 		tc,
 		flagVars,
 		&tmpnet.Network{
-			Owner: "avalanchego-e2e",
+			Owner: flagVars.NetworkOwner(),
 			DefaultFlags: tmpnet.FlagsMap{
 				config.UpgradeFileContentKey: upgradeBase64,
 				// Ensure a min stake duration compatible with testing staking logic

tests/fixture/e2e/env.go

@@ -41,6 +41,8 @@ func InitSharedTestEnvironment(t require.TestingT, envBytes []byte) {
 }
 
 type TestEnvironment struct {
+	// The parent directory of network directories
+	RootNetworkDir string
 	// The directory where the test network configuration is stored
 	NetworkDir string
 	// URIs used to access the API endpoints of nodes of the network
@@ -61,6 +63,7 @@ func GetEnv(tc tests.TestContext) *TestEnvironment {
 		return nil
 	}
 	return &TestEnvironment{
+		RootNetworkDir:              env.RootNetworkDir,
 		NetworkDir:                  env.NetworkDir,
 		URIs:                        env.URIs,
 		PreFundedKey:                env.PreFundedKey,
@@ -81,8 +84,11 @@ func NewTestEnvironment(tc tests.TestContext, flagVars *FlagVars, desiredNetwork
 
 	var network *tmpnet.Network
 
+	networkCmd, err := flagVars.NetworkCmd()
+	require.NoError(err)
+
 	// Consider monitoring flags for any command but stop
-	if !flagVars.StopNetwork() {
+	if networkCmd != StopNetworkCmd {
 		if flagVars.StartCollectors() {
 			require.NoError(tmpnet.StartCollectors(tc.DefaultContext(), tc.Log()))
 		}
@@ -100,8 +106,8 @@ func NewTestEnvironment(tc tests.TestContext, flagVars *FlagVars, desiredNetwork
 		}
 	}
 
-	// Need to load the network if it is being stopped or reused
-	if flagVars.StopNetwork() || flagVars.ReuseNetwork() {
+	// Attempt to load the network if it may already be running
+	if networkCmd == StopNetworkCmd || networkCmd == ReuseNetworkCmd || networkCmd == RestartNetworkCmd {
 		networkDir := flagVars.NetworkDir()
 		var networkSymlink string // If populated, prompts removal of the referenced symlink if --stop-network is specified
 		if len(networkDir) == 0 {
@@ -127,7 +133,7 @@ func NewTestEnvironment(tc tests.TestContext, flagVars *FlagVars, desiredNetwork
 			)
 		}
 
-		if flagVars.StopNetwork() {
+		if networkCmd == StopNetworkCmd {
 			if len(networkSymlink) > 0 {
 				// Remove the symlink to avoid attempts to reuse the stopped network
 				tc.Log().Info("removing symlink",
@@ -144,14 +150,18 @@ func NewTestEnvironment(tc tests.TestContext, flagVars *FlagVars, desiredNetwork
 				tc.Log().Warn("no network to stop")
 			}
 			os.Exit(0)
-		} else if network != nil && flagVars.RestartNetwork() {
-			// A network is only restarted if it is already running and stop was not requested
+		}
+
+		if network != nil && networkCmd == RestartNetworkCmd {
 			require.NoError(network.Restart(tc.DefaultContext(), tc.Log()))
 		}
 	}
 
 	// Start a new network
 	if network == nil {
+		// TODO(marun) Maybe accept a factory function for the desired network
+		// that is only run when a new network will be started?
+
 		network = desiredNetwork
 		runtimeConfig, err := flagVars.NodeRuntimeConfig()
 		require.NoError(err)
@@ -160,9 +170,9 @@ func NewTestEnvironment(tc tests.TestContext, flagVars *FlagVars, desiredNetwork
 		StartNetwork(
 			tc,
 			network,
+			flagVars.RootNetworkDir(),
 			flagVars.NetworkShutdownDelay(),
-			flagVars.StartNetwork(),
-			flagVars.ReuseNetwork(),
+			networkCmd,
 		)
 	}
 
@@ -171,7 +181,7 @@ func NewTestEnvironment(tc tests.TestContext, flagVars *FlagVars, desiredNetwork
 		require.NoError(tmpnet.WaitForPromtailReadiness(tc.DefaultContext(), tc.Log()))
 	}
 
-	if flagVars.StartNetwork() {
+	if networkCmd == StartNetworkCmd {
 		os.Exit(0)
 	}
 
@@ -189,6 +199,7 @@ func NewTestEnvironment(tc tests.TestContext, flagVars *FlagVars, desiredNetwork
 	)
 
 	return &TestEnvironment{
+		RootNetworkDir:              flagVars.RootNetworkDir(),
 		NetworkDir:                  network.Dir,
 		URIs:                        uris,
 		PrivateNetworkShutdownDelay: flagVars.NetworkShutdownDelay(),
@@ -231,8 +242,8 @@ func (te *TestEnvironment) StartPrivateNetwork(network *tmpnet.Network) {
 	StartNetwork(
 		te.testContext,
 		network,
+		te.RootNetworkDir,
 		te.PrivateNetworkShutdownDelay,
-		false, /* skipShutdown */
-		false, /* reuseNetwork */
+		EmptyNetworkCmd,
 	)
 }

tests/fixture/e2e/flags.go

@@ -4,6 +4,7 @@
 package e2e
 
 import (
+	"errors"
 	"flag"
 	"fmt"
 	"os"
@@ -15,39 +16,71 @@ import (
 	"github.com/ava-labs/avalanchego/tests/fixture/tmpnet/flags"
 )
 
+type NetworkCmd int
+
+const (
+	EmptyNetworkCmd NetworkCmd = iota
+	StartNetworkCmd
+	StopNetworkCmd
+	RestartNetworkCmd
+	ReuseNetworkCmd
+)
+
 type FlagVars struct {
-	runtimeConfigVars *flags.RuntimeConfigVars
-	networkDir        string
-	reuseNetwork      bool
-	startCollectors   bool
-	checkMonitoring   bool
-	startNetwork      bool
-	stopNetwork       bool
-	restartNetwork    bool
-	nodeCount         int
-	activateFortuna   bool
-}
+	startNetwork     bool
+	startNetworkVars *flags.StartNetworkVars
 
-func (v *FlagVars) NodeRuntimeConfig() (*tmpnet.NodeRuntimeConfig, error) {
-	return v.runtimeConfigVars.GetNodeRuntimeConfig()
+	startCollectors bool
+	checkMonitoring bool
+
+	networkDir     string
+	reuseNetwork   bool
+	stopNetwork    bool
+	restartNetwork bool
+
+	activateFortuna bool
 }
 
-func (v *FlagVars) NetworkDir() string {
-	if !v.reuseNetwork {
-		return ""
+func (v *FlagVars) NetworkCmd() (NetworkCmd, error) {
+	cmd := EmptyNetworkCmd
+	count := 0
+	if v.startNetwork {
+		cmd = StartNetworkCmd
+		count++
 	}
-	if len(v.networkDir) > 0 {
-		return v.networkDir
+	if v.stopNetwork {
+		cmd = StopNetworkCmd
+		count++
 	}
-	return os.Getenv(tmpnet.NetworkDirEnvName)
+	if v.restartNetwork {
+		cmd = RestartNetworkCmd
+		count++
+	}
+	if v.reuseNetwork {
+		cmd = ReuseNetworkCmd
+		count++
+	}
+	if count > 1 {
+		return EmptyNetworkCmd, errors.New("only one of --start-network, --stop-network, --restart-network, or --reuse-network can be specified")
+	}
+
+	return cmd, nil
 }
 
-func (v *FlagVars) ReuseNetwork() bool {
-	return v.reuseNetwork
+func (v *FlagVars) RootNetworkDir() string {
+	return v.startNetworkVars.RootNetworkDir
 }
 
-func (v *FlagVars) RestartNetwork() bool {
-	return v.restartNetwork
+func (v *FlagVars) NetworkOwner() string {
+	return v.startNetworkVars.NetworkOwner
+}
+
+func (v *FlagVars) NodeCount() (int, error) {
+	return v.startNetworkVars.GetNodeCount()
+}
+
+func (v *FlagVars) NodeRuntimeConfig() (*tmpnet.NodeRuntimeConfig, error) {
+	return v.startNetworkVars.GetNodeRuntimeConfig()
 }
 
 func (v *FlagVars) StartCollectors() bool {
@@ -58,6 +91,16 @@ func (v *FlagVars) CheckMonitoring() bool {
 	return v.checkMonitoring
 }
 
+func (v *FlagVars) NetworkDir() string {
+	if !v.reuseNetwork {
+		return ""
+	}
+	if len(v.networkDir) > 0 {
+		return v.networkDir
+	}
+	return os.Getenv(tmpnet.NetworkDirEnvName)
+}
+
 func (v *FlagVars) NetworkShutdownDelay() time.Duration {
 	if v.startCollectors {
 		// Only return a non-zero value if we want to ensure the collectors have
@@ -67,65 +110,59 @@ func (v *FlagVars) NetworkShutdownDelay() time.Duration {
 	return 0
 }
 
-func (v *FlagVars) StartNetwork() bool {
-	return v.startNetwork
-}
-
-func (v *FlagVars) StopNetwork() bool {
-	return v.stopNetwork
-}
-
-func (v *FlagVars) NodeCount() int {
-	return v.nodeCount
-}
-
 func (v *FlagVars) ActivateFortuna() bool {
 	return v.activateFortuna
 }
 
 func RegisterFlags() *FlagVars {
+	return RegisterFlagsWithDefaultOwner("")
+}
+
+func RegisterFlagsWithDefaultOwner(defaultOwner string) *FlagVars {
 	vars := FlagVars{}
-	vars.runtimeConfigVars = flags.NewRuntimeConfigFlagVars()
+
+	flag.BoolVar(
+		&vars.startNetwork,
+		"start-network",
+		false,
+		"[optional] start a new network and exit without executing any tests. The new network cannot be reused with --reuse-network.",
+	)
+
+	vars.startNetworkVars = flags.NewStartNetworkFlagVars(defaultOwner)
+
+	SetMonitoringFlags(
+		&vars.startCollectors,
+		&vars.checkMonitoring,
+	)
+
 	flag.StringVar(
 		&vars.networkDir,
 		"network-dir",
-		"",
-		fmt.Sprintf("[optional] the dir containing the configuration of an existing network to target for testing. Will only be used if --reuse-network is specified. Also possible to configure via the %s env variable.", tmpnet.NetworkDirEnvName),
+		tmpnet.GetEnvWithDefault(tmpnet.NetworkDirEnvName, ""),
+		fmt.Sprintf("[optional] the dir containing the configuration of an existing network. Will only be used if --reuse-network, --restart-network or --stop-network are specified. Also possible to configure via the %s env variable.", tmpnet.NetworkDirEnvName),
 	)
+
 	flag.BoolVar(
 		&vars.reuseNetwork,
 		"reuse-network",
 		false,
-		"[optional] reuse an existing network previously started with --reuse-network. If a network is not already running, create a new one and leave it running for subsequent usage. Ignored if --stop-network is provided.",
+		"[optional] run tests against an existing network previously started with --reuse-network. If a network is not already running, create a new one and leave it running for subsequent usage.",
 	)
+
 	flag.BoolVar(
 		&vars.restartNetwork,
 		"restart-network",
 		false,
-		"[optional] restart an existing network previously started with --reuse-network. Useful for ensuring a network is running with the current state of binaries on disk. Ignored if a network is not already running or --stop-network is provided.",
-	)
-	SetMonitoringFlags(
-		&vars.startCollectors,
-		&vars.checkMonitoring,
-	)
-	flag.BoolVar(
-		&vars.startNetwork,
-		"start-network",
-		false,
-		"[optional] start a new network and exit without executing any tests. The new network cannot be reused with --reuse-network. Ignored if either --reuse-network or --stop-network is provided.",
+		"[optional] like --reuse-network except an already running network is restarted before running tests to ensure the network represents the current state of binaries on disk.",
 	)
+
 	flag.BoolVar(
 		&vars.stopNetwork,
 		"stop-network",
 		false,
 		"[optional] stop an existing network started with --reuse-network and exit without executing any tests.",
 	)
-	flag.IntVar(
-		&vars.nodeCount,
-		"node-count",
-		tmpnet.DefaultNodeCount,
-		"number of nodes the network should initially consist of",
-	)
+
 	flag.BoolVar(
 		&vars.activateFortuna,
 		"activate-fortuna",