github.com/badrootd/celestia-core@v0.0.0-20240305091328-aa4207a4b25d/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/badrootd/celestia-core/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 -d dir [-g int] [-m version_weight_csv] [-p]",
    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  			multiVersion, err := cmd.Flags().GetString("multi-version")
    48  			if err != nil {
    49  				return err
    50  			}
    51  			prometheus, err := cmd.Flags().GetBool("prometheus")
    52  			if err != nil {
    53  				return err
    54  			}
    55  			return cli.generate(dir, groups, multiVersion, prometheus)
    56  		},
    57  	}
    58  
    59  	cli.root.PersistentFlags().StringP("dir", "d", "", "Output directory for manifests")
    60  	_ = cli.root.MarkPersistentFlagRequired("dir")
    61  	cli.root.PersistentFlags().StringP("multi-version", "m", "", "Comma-separated list of versions of CometBFT to test in the generated testnets, "+
    62  		"or empty to only use this branch's version")
    63  	cli.root.PersistentFlags().IntP("groups", "g", 0, "Number of groups")
    64  	cli.root.PersistentFlags().BoolP("prometheus", "p", false, "Enable generation of Prometheus metrics on all manifests")
    65  
    66  	return cli
    67  }
    68  
    69  // generate generates manifests in a directory.
    70  func (cli *CLI) generate(dir string, groups int, multiVersion string, prometheus bool) error {
    71  	err := os.MkdirAll(dir, 0o755)
    72  	if err != nil {
    73  		return err
    74  	}
    75  
    76  	cfg := &generateConfig{
    77  		randSource:   rand.New(rand.NewSource(randomSeed)), //nolint:gosec
    78  		multiVersion: multiVersion,
    79  		prometheus:   prometheus,
    80  	}
    81  	manifests, err := Generate(cfg)
    82  	if err != nil {
    83  		return err
    84  	}
    85  	if groups <= 0 {
    86  		for i, manifest := range manifests {
    87  			err = manifest.Save(filepath.Join(dir, fmt.Sprintf("gen-%04d.toml", i)))
    88  			if err != nil {
    89  				return err
    90  			}
    91  		}
    92  	} else {
    93  		groupSize := int(math.Ceil(float64(len(manifests)) / float64(groups)))
    94  		for g := 0; g < groups; g++ {
    95  			for i := 0; i < groupSize && g*groupSize+i < len(manifests); i++ {
    96  				manifest := manifests[g*groupSize+i]
    97  				err = manifest.Save(filepath.Join(dir, fmt.Sprintf("gen-group%02d-%04d.toml", g, i)))
    98  				if err != nil {
    99  					return err
   100  				}
   101  			}
   102  		}
   103  	}
   104  	return nil
   105  }
   106  
   107  // Run runs the CLI.
   108  func (cli *CLI) Run() {
   109  	if err := cli.root.Execute(); err != nil {
   110  		logger.Error(err.Error())
   111  		os.Exit(1)
   112  	}
   113  }