github.com/MetalBlockchain/metalgo@v1.11.9/tests/antithesis/config.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package antithesis
     5  
     6  import (
     7  	"errors"
     8  	"fmt"
     9  
    10  	"github.com/spf13/pflag"
    11  	"github.com/spf13/viper"
    12  
    13  	"github.com/MetalBlockchain/metalgo/config"
    14  	"github.com/MetalBlockchain/metalgo/wallet/subnet/primary"
    15  )
    16  
    17  const (
    18  	URIsKey     = "uris"
    19  	ChainIDsKey = "chain-ids"
    20  
    21  	FlagsName = "workload"
    22  	EnvPrefix = "avawl"
    23  )
    24  
    25  var (
    26  	errNoURIs      = errors.New("at least one URI must be provided")
    27  	errNoArguments = errors.New("no arguments")
    28  )
    29  
    30  type Config struct {
    31  	URIs     []string
    32  	ChainIDs []string
    33  }
    34  
    35  func NewConfig(arguments []string) (*Config, error) {
    36  	v, err := parseFlags(arguments)
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  
    41  	c := &Config{
    42  		URIs:     v.GetStringSlice(URIsKey),
    43  		ChainIDs: v.GetStringSlice(ChainIDsKey),
    44  	}
    45  	return c, c.Verify()
    46  }
    47  
    48  func (c *Config) Verify() error {
    49  	if len(c.URIs) == 0 {
    50  		return errNoURIs
    51  	}
    52  	return nil
    53  }
    54  
    55  func parseFlags(arguments []string) (*viper.Viper, error) {
    56  	if len(arguments) == 0 {
    57  		return nil, errNoArguments
    58  	}
    59  
    60  	fs := pflag.NewFlagSet(FlagsName, pflag.ContinueOnError)
    61  	fs.StringSlice(URIsKey, []string{primary.LocalAPIURI}, "URIs of nodes that the workload can communicate with")
    62  	fs.StringSlice(ChainIDsKey, []string{}, "IDs of chains to target for testing")
    63  	if err := fs.Parse(arguments[1:]); err != nil {
    64  		return nil, fmt.Errorf("failed parsing CLI flags: %w", err)
    65  	}
    66  
    67  	v := viper.New()
    68  	v.AutomaticEnv()
    69  	v.SetEnvKeyReplacer(config.DashesToUnderscores)
    70  	v.SetEnvPrefix(EnvPrefix)
    71  	if err := v.BindPFlags(fs); err != nil {
    72  		return nil, fmt.Errorf("failed binding pflags: %w", err)
    73  	}
    74  	return v, nil
    75  }