github.com/yandex/pandora@v0.5.32/components/providers/scenario/config/decode.go (about)

     1  package config
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"strconv"
     7  
     8  	"github.com/yandex/pandora/components/providers/scenario/vs"
     9  	"github.com/yandex/pandora/core/config"
    10  	"github.com/yandex/pandora/lib/math"
    11  	"github.com/yandex/pandora/lib/str"
    12  	"gopkg.in/yaml.v2"
    13  )
    14  
    15  func ParseAmmoConfig(file io.Reader) (*AmmoConfig, error) {
    16  	const op = "scenario/decoder.ParseAmmoConfig"
    17  	bytes, err := io.ReadAll(file)
    18  	if err != nil {
    19  		return nil, fmt.Errorf("%s, io.ReadAll, %w", op, err)
    20  	}
    21  	cfg, err := DecodeMap(bytes)
    22  	if err != nil {
    23  		return nil, fmt.Errorf("%s, decodeMap, %w", op, err)
    24  	}
    25  	return cfg, nil
    26  }
    27  
    28  func ConvertHCLToAmmo(ammo AmmoHCL) (*AmmoConfig, error) {
    29  	const op = "scenario.ConvertHCLToAmmo"
    30  	bytes, err := yaml.Marshal(ammo)
    31  	if err != nil {
    32  		return nil, fmt.Errorf("%s, cant yaml.Marshal: %w", op, err)
    33  	}
    34  	cfg, err := DecodeMap(bytes)
    35  	if err != nil {
    36  		return nil, fmt.Errorf("%s, decodeMap, %w", op, err)
    37  	}
    38  	return cfg, nil
    39  }
    40  
    41  func DecodeMap(bytes []byte) (*AmmoConfig, error) {
    42  	const op = "scenario/decoder.decodeMap"
    43  
    44  	var ammoCfg AmmoConfig
    45  
    46  	data := make(map[string]any)
    47  	err := yaml.Unmarshal(bytes, &data)
    48  	if err != nil {
    49  		return nil, fmt.Errorf("%s, yaml.Unmarshal, %w", op, err)
    50  	}
    51  	err = config.DecodeAndValidate(data, &ammoCfg)
    52  	if err != nil {
    53  		return nil, fmt.Errorf("%s, config.DecodeAndValidate, %w", op, err)
    54  	}
    55  	return &ammoCfg, nil
    56  }
    57  
    58  func ExtractVariableStorage(cfg *AmmoConfig) (*vs.SourceStorage, error) {
    59  	storage := vs.NewVariableStorage()
    60  	for _, source := range cfg.VariableSources {
    61  		err := source.Init()
    62  		if err != nil {
    63  			return storage, err
    64  		}
    65  		storage.AddSource(source.GetName(), source.GetVariables())
    66  	}
    67  	return storage, nil
    68  }
    69  
    70  func ParseShootName(shoot string) (string, int, int, error) {
    71  	name, args, err := str.ParseStringFunc(shoot)
    72  	if err != nil {
    73  		return "", 0, 0, err
    74  	}
    75  	cnt := 1
    76  	if len(args) > 0 && args[0] != "" {
    77  		cnt, err = strconv.Atoi(args[0])
    78  		if err != nil {
    79  			return "", 0, 0, fmt.Errorf("failed to parse count: %w", err)
    80  		}
    81  	}
    82  	sleep := 0
    83  	if len(args) > 1 && args[1] != "" {
    84  		sleep, err = strconv.Atoi(args[1])
    85  		if err != nil {
    86  			return "", 0, 0, fmt.Errorf("failed to parse count: %w", err)
    87  		}
    88  	}
    89  	return name, cnt, sleep, nil
    90  }
    91  
    92  func SpreadNames(input []ScenarioConfig) (map[string]int, int) {
    93  	if len(input) == 0 {
    94  		return nil, 0
    95  	}
    96  	if len(input) == 1 {
    97  		return map[string]int{input[0].Name: 1}, 1
    98  	}
    99  
   100  	scenarioRegistry := map[string]ScenarioConfig{}
   101  	weights := make([]int64, len(input))
   102  	for i := range input {
   103  		scenarioRegistry[input[i].Name] = input[i]
   104  		if input[i].Weight == 0 {
   105  			input[i].Weight = 1
   106  		}
   107  		weights[i] = input[i].Weight
   108  	}
   109  
   110  	div := math.GCDM(weights...)
   111  	names := make(map[string]int)
   112  	total := 0
   113  	for _, sc := range input {
   114  		cnt := int(sc.Weight / div)
   115  		total += cnt
   116  		names[sc.Name] = cnt
   117  	}
   118  	return names, total
   119  }