github.com/vipernet-xyz/tm@v0.34.24/test/e2e/generator/main.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"math"
     6  	"math/rand"
     7  	"os"
     8  	"path/filepath"
     9  
    10  	"github.com/spf13/cobra"
    11  
    12  	"github.com/vipernet-xyz/tm/libs/log"
    13  )
    14  
    15  const (
    16  	randomSeed int64 = 4827085738
    17  )
    18  
    19  var logger = log.NewTMLogger(log.NewSyncWriter(os.Stdout))
    20  
    21  func main() {
    22  	NewCLI().Run()
    23  }
    24  
    25  // CLI is the Cobra-based command-line interface.
    26  type CLI struct {
    27  	root *cobra.Command
    28  }
    29  
    30  // NewCLI sets up the CLI.
    31  func NewCLI() *CLI {
    32  	cli := &CLI{}
    33  	cli.root = &cobra.Command{
    34  		Use:           "generator",
    35  		Short:         "End-to-end testnet generator",
    36  		SilenceUsage:  true,
    37  		SilenceErrors: true, // we'll output them ourselves in Run()
    38  		RunE: func(cmd *cobra.Command, args []string) error {
    39  			dir, err := cmd.Flags().GetString("dir")
    40  			if err != nil {
    41  				return err
    42  			}
    43  			groups, err := cmd.Flags().GetInt("groups")
    44  			if err != nil {
    45  				return err
    46  			}
    47  			return cli.generate(dir, groups)
    48  		},
    49  	}
    50  
    51  	cli.root.PersistentFlags().StringP("dir", "d", "", "Output directory for manifests")
    52  	_ = cli.root.MarkPersistentFlagRequired("dir")
    53  	cli.root.PersistentFlags().IntP("groups", "g", 0, "Number of groups")
    54  
    55  	return cli
    56  }
    57  
    58  // generate generates manifests in a directory.
    59  func (cli *CLI) generate(dir string, groups int) error {
    60  	err := os.MkdirAll(dir, 0o755)
    61  	if err != nil {
    62  		return err
    63  	}
    64  
    65  	//nolint:gosec // G404: Use of weak random number generator (math/rand instead of crypto/rand)
    66  	manifests, err := Generate(rand.New(rand.NewSource(randomSeed)))
    67  	if err != nil {
    68  		return err
    69  	}
    70  	if groups <= 0 {
    71  		for i, manifest := range manifests {
    72  			err = manifest.Save(filepath.Join(dir, fmt.Sprintf("gen-%04d.toml", i)))
    73  			if err != nil {
    74  				return err
    75  			}
    76  		}
    77  	} else {
    78  		groupSize := int(math.Ceil(float64(len(manifests)) / float64(groups)))
    79  		for g := 0; g < groups; g++ {
    80  			for i := 0; i < groupSize && g*groupSize+i < len(manifests); i++ {
    81  				manifest := manifests[g*groupSize+i]
    82  				err = manifest.Save(filepath.Join(dir, fmt.Sprintf("gen-group%02d-%04d.toml", g, i)))
    83  				if err != nil {
    84  					return err
    85  				}
    86  			}
    87  		}
    88  	}
    89  	return nil
    90  }
    91  
    92  // Run runs the CLI.
    93  func (cli *CLI) Run() {
    94  	if err := cli.root.Execute(); err != nil {
    95  		logger.Error(err.Error())
    96  		os.Exit(1)
    97  	}
    98  }