feat: improve set-env CLI cmd (#6001)

* improve set-env CLI cmd

* update changelog

* add optional chain id to localnet

---------

Co-authored-by: devbot-wizard <[email protected]>

CHANGELOG.md

@@ -85,6 +85,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#5890](https://github.com/osmosis-labs/osmosis/pull/5890) feat: CreateCLPool & LinkCFMMtoCL pool into one gov-prop
 * [#5964](https://github.com/osmosis-labs/osmosis/pull/5964) fix e2e test concurrency bugs
 * [#5948] (https://github.com/osmosis-labs/osmosis/pull/5948) Parameterizing Pool Type Information in Protorev
+* [#6001](https://github.com/osmosis-labs/osmosis/pull/6001) feat: improve set-env CLI cmd
 
 ### Minor improvements & Bug Fixes
 

cmd/osmosisd/cmd/change_environment.go

@@ -1,10 +1,13 @@
 package cmd
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
 
+	"github.com/cosmos/cosmos-sdk/client"
+	"github.com/cosmos/cosmos-sdk/client/flags"
 	"github.com/joho/godotenv"
 	"github.com/spf13/cobra"
 
@@ -14,7 +17,8 @@ import (
 const (
 	EnvVariable = "OSMOSISD_ENVIRONMENT"
 	EnvMainnet  = "mainnet"
-	EnvLocalnet = "localosmosis"
+	EnvTestnet  = "testnet"
+	EnvLocalnet = "localnet"
 )
 
 // ExportAirdropSnapshotCmd generates a snapshot.json from a provided exported genesis.json.
@@ -25,29 +29,22 @@ func ChangeEnvironmentCmd() *cobra.Command {
 		Long: `Set home environment variables for commands
 Example:
 	osmosisd set-env mainnet
-	osmosisd set-env localosmosis
+	osmosisd set-env testnet
+	osmosisd set-env localnet [optional-chain-id]
 	osmosisd set-env $HOME/.custom-dir
 `,
-		Args: cobra.ExactArgs(1),
+		Args: customArgs,
 		RunE: func(cmd *cobra.Command, args []string) error {
+			// Note: If we are calling this method, the environment file has already been set in
+			// NewRootCmd() when creating the rootCmd. We do this because order of operations
+			// dictates this as a requirement. If we changed the env file here, the osmosis
+			// daemon would not initialize the folder we are intending to set to.
 			newEnv := args[0]
-
-			currentEnvironment := getHomeEnvironment()
-			fmt.Println("Current environment: ", currentEnvironment)
-
-			if _, err := environmentNameToPath(newEnv); err != nil {
-				return err
-			}
-
-			fmt.Println("New environment: ", newEnv)
-
-			envMap := make(map[string]string)
-			envMap[EnvVariable] = newEnv
-			err := godotenv.Write(envMap, filepath.Join(app.DefaultNodeHome, ".env"))
-			if err != nil {
-				return err
+			chainId := ""
+			if len(args) > 1 {
+				chainId = args[1]
 			}
-			return nil
+			return clientSettingsFromEnv(cmd, newEnv, chainId)
 		},
 	}
 	return cmd
@@ -64,8 +61,8 @@ Example:
 
 	Returns one of:
 	- mainnet implying $HOME/.osmosisd
+	- testnet implying $HOME/.osmosisd-test
 	- localosmosis implying $HOME/.osmosisd-local
-	- localosmosis
 	- custom path`,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			environment := getHomeEnvironment()
@@ -91,17 +88,141 @@ func environmentNameToPath(environmentName string) (string, error) {
 	switch environmentName {
 	case EnvMainnet:
 		return app.DefaultNodeHome, nil
+	case EnvTestnet:
+		return filepath.Join(userHomeDir, ".osmosisd-test/"), nil
 	case EnvLocalnet:
 		return filepath.Join(userHomeDir, ".osmosisd-local/"), nil
 	default:
-		osmosisdPath := filepath.Join(userHomeDir, environmentName)
-		_, err := os.Stat(osmosisdPath)
+		_, err := os.Stat(environmentName)
 		if os.IsNotExist(err) {
 			// Creating new environment directory
-			if err := os.Mkdir(osmosisdPath, os.ModePerm); err != nil {
+			if err := os.Mkdir(environmentName, os.ModePerm); err != nil {
 				return "", err
 			}
 		}
-		return osmosisdPath, nil
+		return environmentName, nil
+	}
+}
+
+// clientSettingsFromEnv takes the env name (mainnet, testnet, localnet, etc) and sets the
+// client.toml settings to commonly used values for that environment.
+func clientSettingsFromEnv(cmd *cobra.Command, environmentName, chainId string) error {
+	envConfigs := map[string]map[string]string{
+		EnvMainnet: {
+			flags.FlagChainID:       "osmosis-1",
+			flags.FlagNode:          "https://rpc.osmosis.zone:443",
+			flags.FlagBroadcastMode: "block",
+		},
+		EnvTestnet: {
+			flags.FlagChainID:       "osmo-test-5",
+			flags.FlagNode:          "https://rpc.testnet.osmosis.zone:443",
+			flags.FlagBroadcastMode: "block",
+		},
+		EnvLocalnet: {
+			flags.FlagChainID:       "localosmosis",
+			flags.FlagBroadcastMode: "block",
+		},
+	}
+
+	configs, ok := envConfigs[environmentName]
+	if !ok {
+		return nil
+	}
+
+	// Update the ChainID if environmentName is EnvLocalnet and chainId is provided
+	if environmentName == EnvLocalnet && chainId != "" {
+		configs[flags.FlagChainID] = chainId
+	}
+
+	for flag, value := range configs {
+		if err := runConfigCmd(cmd, []string{flag, value}); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// changeEnvironment takes the given environment name and changes the .env file to reflect it.
+func changeEnvironment(args []string) error {
+	newEnv := args[0]
+
+	currentEnvironment := getHomeEnvironment()
+	fmt.Println("Current environment: ", currentEnvironment)
+
+	if _, err := environmentNameToPath(newEnv); err != nil {
+		return err
+	}
+
+	fmt.Println("New environment: ", newEnv)
+
+	envMap := make(map[string]string)
+	envMap[EnvVariable] = newEnv
+	err := godotenv.Write(envMap, filepath.Join(app.DefaultNodeHome, ".env"))
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// createHomeDirIfNotExist creates the home directory if it does not exist and writes a blank
+// .env file. This is used for the first time setup of the osmosisd home directory.
+func createHomeDirIfNotExist(homeDir string) error {
+	if _, err := os.Stat(homeDir); os.IsNotExist(err) {
+		err := os.MkdirAll(homeDir, 0755)
+		if err != nil {
+			return err
+		}
+	}
+
+	envFilePath := filepath.Join(homeDir, ".env")
+	if _, err := os.Stat(envFilePath); os.IsNotExist(err) {
+		file, err := os.Create(envFilePath)
+		if err != nil {
+			return err
+		}
+		file.Close()
+	}
+
+	return nil
+}
+
+// changeEnvPriorToSetup changes the env file to reflect the desired environment the user wants to change to.
+// If this is not called in NewRootCmd(), the environment change will happen **after** all relevant setup actions
+// happen (e.g., the .env will be read in as the previous value in the setup and not the new value).
+func changeEnvPriorToSetup(cmd *cobra.Command, initClientCtx *client.Context, args []string, homeDir string) error {
+	if cmd.Name() == "set-env" {
+		err := createHomeDirIfNotExist(homeDir)
+		if err != nil {
+			return err
+		}
+
+		err = changeEnvironment(args)
+		if err != nil {
+			return err
+		}
+
+		homeEnvironment := getHomeEnvironment()
+		homeDir, err := environmentNameToPath(homeEnvironment)
+		if err != nil {
+			// Failed to convert home environment to home path, using default home
+			homeDir = app.DefaultNodeHome
+		}
+		*initClientCtx = initClientCtx.WithHomeDir(homeDir)
+	}
+	return nil
+}
+
+// customArgs accepts one arg, but if the first arg is "localnet", then it accepts two args.
+func customArgs(cmd *cobra.Command, args []string) error {
+	if len(args) < 1 || len(args) > 2 {
+		return errors.New("set-env requires 1 or 2 arguments")
+	}
+	if args[0] == "localnet" {
+		return nil
+	}
+	if len(args) == 2 {
+		return errors.New("only 'set-env localnet' accepts a second argument")
 	}
+	return nil
 }

cmd/osmosisd/cmd/root.go

@@ -287,6 +287,12 @@ func NewRootCmd() (*cobra.Command, params.EncodingConfig) {
 		Use:   "osmosisd",
 		Short: "Start osmosis app",
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
+			// If not calling the set-env command, this is a no-op.
+			err := changeEnvPriorToSetup(cmd, &initClientCtx, args, homeDir)
+			if err != nil {
+				return err
+			}
+
 			initClientCtx, err := client.ReadPersistentCommandFlags(initClientCtx, cmd.Flags())
 			if err != nil {
 				return err
@@ -391,8 +397,9 @@ func NewRootCmd() (*cobra.Command, params.EncodingConfig) {
 func getHomeEnvironment() string {
 	envPath := filepath.Join(osmosis.DefaultNodeHome, ".env")
 
-	// Use default node home if can't get environment
-	err := godotenv.Load(envPath)
+	// Use default node home if can't get environment.
+	// Overload must be used here in the event that the .env gets updated.
+	err := godotenv.Overload(envPath)
 	if err != nil {
 		// Failed to load, using default home directory
 		return EnvMainnet

tests/e2e/e2e_test.go

@@ -3,7 +3,6 @@ package e2e
 import (
 	"encoding/json"
 	"fmt"
-	"github.com/cosmos/cosmos-sdk/types/address"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -12,6 +11,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/cosmos/cosmos-sdk/types/address"
+
 	transfertypes "github.com/cosmos/ibc-go/v4/modules/apps/transfer/types"
 	"github.com/iancoleman/orderedmap"