[tmpnet] Unify start network flag usage between e2e and tmpnetctl

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,
@@ -152,6 +155,9 @@ func NewTestEnvironment(tc tests.TestContext, flagVars *FlagVars, desiredNetwork
 
 	// 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,6 +166,7 @@ func NewTestEnvironment(tc tests.TestContext, flagVars *FlagVars, desiredNetwork
 		StartNetwork(
 			tc,
 			network,
+			flagVars.RootNetworkDir(),
 			flagVars.NetworkShutdownDelay(),
 			flagVars.StartNetwork(),
 			flagVars.ReuseNetwork(),
@@ -189,6 +196,7 @@ func NewTestEnvironment(tc tests.TestContext, flagVars *FlagVars, desiredNetwork
 	)
 
 	return &TestEnvironment{
+		RootNetworkDir:              flagVars.RootNetworkDir(),
 		NetworkDir:                  network.Dir,
 		URIs:                        uris,
 		PrivateNetworkShutdownDelay: flagVars.NetworkShutdownDelay(),
@@ -231,6 +239,7 @@ func (te *TestEnvironment) StartPrivateNetwork(network *tmpnet.Network) {
 	StartNetwork(
 		te.testContext,
 		network,
+		te.RootNetworkDir,
 		te.PrivateNetworkShutdownDelay,
 		false, /* skipShutdown */
 		false, /* reuseNetwork */

tests/fixture/e2e/flags.go

@@ -16,20 +16,46 @@ import (
 )
 
 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
+
+	startCollectors bool
+	checkMonitoring bool
+
+	networkDir     string
+	reuseNetwork   bool
+	restartNetwork bool
+	stopNetwork    bool
+
+	activateFortuna bool
+}
+
+func (v *FlagVars) StartNetwork() bool {
+	return v.startNetwork
+}
+
+func (v *FlagVars) RootNetworkDir() string {
+	return v.startNetworkVars.RootNetworkDir
+}
+
+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.runtimeConfigVars.GetNodeRuntimeConfig()
+	return v.startNetworkVars.GetNodeRuntimeConfig()
+}
+
+func (v *FlagVars) StartCollectors() bool {
+	return v.startCollectors
+}
+
+func (v *FlagVars) CheckMonitoring() bool {
+	return v.checkMonitoring
 }
 
 func (v *FlagVars) NetworkDir() string {
@@ -50,12 +76,8 @@ func (v *FlagVars) RestartNetwork() bool {
 	return v.restartNetwork
 }
 
-func (v *FlagVars) StartCollectors() bool {
-	return v.startCollectors
-}
-
-func (v *FlagVars) CheckMonitoring() bool {
-	return v.checkMonitoring
+func (v *FlagVars) StopNetwork() bool {
+	return v.stopNetwork
 }
 
 func (v *FlagVars) NetworkShutdownDelay() time.Duration {
@@ -67,65 +89,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. Ignored if either --reuse-network or --stop-network is provided.",
+	)
+
+	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),
+		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.",
 	)
+
 	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.",
-	)
+
 	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",

tests/fixture/e2e/helpers.go

@@ -46,10 +46,6 @@ const (
 
 	DefaultGasLimit = uint64(21000) // Standard gas limit
 
-	// An empty string prompts the use of the default path which ensures a
-	// predictable target for github's upload-artifact action.
-	DefaultNetworkDir = ""
-
 	// Directory used to store private networks (specific to a single test)
 	// under the shared network dir.
 	PrivateNetworksDirName = "private_networks"
@@ -272,6 +268,7 @@ func CheckBootstrapIsPossible(tc tests.TestContext, network *tmpnet.Network) *tm
 func StartNetwork(
 	tc tests.TestContext,
 	network *tmpnet.Network,
+	rootNetworkDir string,
 	shutdownDelay time.Duration,
 	skipShutdown bool,
 	reuseNetwork bool,
@@ -282,7 +279,7 @@ func StartNetwork(
 		tc.DefaultContext(),
 		tc.Log(),
 		network,
-		DefaultNetworkDir,
+		rootNetworkDir,
 	)
 	if err != nil {
 		tc.DeferCleanup(func() {

tests/fixture/tmpnet/README.md

@@ -63,6 +63,7 @@ the following non-test files:
 | flags/common.go             |             | Defines type definitions common across other files                     |
 | flags/process_runtime.go    |             | Defines flags configuring the process node runtime                     |
 | flags/runtime.go            |             | Defines flags configuring node runtime                                 |
+| flags/start_network.go      |             | Defines flags configuring network start                                |
 | tmpnetctl/                  |             | Directory containing main entrypoint for tmpnetctl command             |
 | check_monitoring.go         |             | Enables checking if logs and metrics were collected                    |
 | defaults.go                 |             | Defines common default configuration                                   |
@@ -269,12 +270,12 @@ this file (i.e. `source network.env`) in a shell will configure ginkgo
 e2e and the `tmpnetctl` cli to target the network path specified in
 the env var.
 
-Set `TMPNET_ROOT_DIR` to specify the root directory in which to create
-the configuration directory of new networks
-(e.g. `$TMPNET_ROOT_DIR/[network-dir]`). The default root directory is
-`~/.tmpdir/networks`. Configuring the root directory is only relevant
-when creating new networks as the path of existing networks will
-already have been set.
+Set `TMPNET_ROOT_NETWORK_DIR` to specify the root network directory in
+which to create the configuration directory of new networks
+(e.g. `TMPNET_ROOT_NETWORK_DIR/[network-dir]`). The default network
+root directory is `~/.tmpdir/networks`. Configuring the network root
+directory is only relevant when creating new networks as the path of
+existing networks will already have been set.
 
 ### Node configuration
 [Top](#table-of-contents)

tests/fixture/tmpnet/flags/common.go

@@ -10,3 +10,5 @@ package flags
 type stringVarFunc func(p *string, name string, value string, usage string)
 
 type boolVarFunc func(p *bool, name string, value bool, usage string)
+
+type intVarFunc func(p *int, name string, value int, usage string)

tests/fixture/tmpnet/flags/start_network.go

@@ -0,0 +1,71 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package flags
+
+import (
+	"flag"
+	"fmt"
+
+	"github.com/spf13/pflag"
+
+	"github.com/ava-labs/avalanchego/tests/fixture/tmpnet"
+)
+
+type StartNetworkVars struct {
+	// Accessible directly
+	RootNetworkDir string
+	NetworkOwner   string
+
+	// Accessible via a validating method
+	nodeCount   int
+	runtimeVars *RuntimeConfigVars
+
+	defaultNetworkOwner string
+}
+
+func NewStartNetworkFlagVars(defaultNetworkOwner string) *StartNetworkVars {
+	v := &StartNetworkVars{defaultNetworkOwner: defaultNetworkOwner}
+	v.runtimeVars = NewRuntimeConfigFlagVars()
+	v.register(flag.StringVar, flag.IntVar)
+	return v
+}
+
+func NewStartNetworkFlagSetVars(flagSet *pflag.FlagSet, defaultNetworkOwner string) *StartNetworkVars {
+	v := &StartNetworkVars{defaultNetworkOwner: defaultNetworkOwner}
+	v.runtimeVars = NewRuntimeConfigFlagSetVars(flagSet)
+	v.register(flagSet.StringVar, flagSet.IntVar)
+	return v
+}
+
+func (v *StartNetworkVars) register(stringVar stringVarFunc, intVar intVarFunc) {
+	stringVar(
+		&v.RootNetworkDir,
+		"root-network-dir",
+		// An empty string prompts the use of the default path.
+		tmpnet.GetEnvWithDefault(tmpnet.RootNetworkDirEnvName, ""),
+		fmt.Sprintf("The directory in which to create the network directory. Also possible to configure via the %s env variable.", tmpnet.RootNetworkDirEnvName),
+	)
+	stringVar(
+		&v.NetworkOwner,
+		"network-owner",
+		v.defaultNetworkOwner,
+		"The string identifying the intended owner of the network")
+	intVar(
+		&v.nodeCount,
+		"node-count",
+		tmpnet.DefaultNodeCount,
+		"Number of nodes the network should initially consist of",
+	)
+}
+
+func (v *StartNetworkVars) GetNodeCount() (int, error) {
+	if v.nodeCount < 1 {
+		return 0, fmt.Errorf("--node-count must be greater than 0 but got %d", v.nodeCount)
+	}
+	return v.nodeCount, nil
+}
+
+func (v *StartNetworkVars) GetNodeRuntimeConfig() (*tmpnet.NodeRuntimeConfig, error) {
+	return v.runtimeVars.GetNodeRuntimeConfig()
+}

tests/fixture/tmpnet/network.go

@@ -42,8 +42,8 @@ import (
 const (
 	// Constants defining the names of shell variables whose value can
 	// configure network orchestration.
-	NetworkDirEnvName = "TMPNET_NETWORK_DIR"
-	RootDirEnvName    = "TMPNET_ROOT_DIR"
+	RootNetworkDirEnvName = "TMPNET_ROOT_NETWORK_DIR"
+	NetworkDirEnvName     = "TMPNET_NETWORK_DIR"
 
 	// Message to log indicating where to look for metrics and logs for network
 	MetricsAvailableMessage = "metrics and logs available via grafana (collectors must be running)"