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

     1  package templater
     2  
     3  import (
     4  	"fmt"
     5  	"math/rand"
     6  	"strconv"
     7  	"strings"
     8  	"text/template"
     9  	"time"
    10  
    11  	"github.com/gofrs/uuid"
    12  	"github.com/yandex/pandora/lib/numbers"
    13  	"github.com/yandex/pandora/lib/str"
    14  )
    15  
    16  const defaultMaxRandValue = 10
    17  
    18  func init() {
    19  	rand.New(rand.NewSource(time.Now().UnixNano()))
    20  }
    21  
    22  type templateFunc func(args ...any) (string, error)
    23  
    24  var _ = []templateFunc{
    25  	RandInt,
    26  	RandString,
    27  	UUID,
    28  }
    29  
    30  func ParseFunc(v string) (f any, args []string) {
    31  	name, args := parseStr(v)
    32  	if f, ok := GetFuncs()[name]; ok {
    33  		return f, args
    34  	}
    35  	return nil, nil
    36  }
    37  
    38  func parseStr(v string) (string, []string) {
    39  	args := strings.Split(v, "(")
    40  	if len(args) == 0 {
    41  		return v, nil
    42  	}
    43  	name := args[0]
    44  	if len(args) == 1 {
    45  		return name, nil
    46  	}
    47  	v = strings.TrimSuffix(strings.Join(args[1:], "("), ")")
    48  	args = strings.Split(v, ",")
    49  	if len(args) == 1 && args[0] == "" {
    50  		return name, nil
    51  	}
    52  	for i := 0; i < len(args); i++ {
    53  		args[i] = strings.TrimSpace(args[i])
    54  	}
    55  
    56  	return name, args
    57  }
    58  
    59  func GetFuncs() template.FuncMap {
    60  	return map[string]any{
    61  		"randInt":    RandInt,
    62  		"randString": RandString,
    63  		"uuid":       UUID,
    64  	}
    65  }
    66  
    67  func RandInt(args ...any) (string, error) {
    68  	switch len(args) {
    69  	case 0:
    70  		return randInt(0, 0)
    71  	case 1:
    72  		f, err := numbers.ParseInt(args[0])
    73  		if err != nil {
    74  			return "", err
    75  		}
    76  		return randInt(f, 0)
    77  	case 2:
    78  		f, err := numbers.ParseInt(args[0])
    79  		if err != nil {
    80  			return "", err
    81  		}
    82  		t, err := numbers.ParseInt(args[1])
    83  		if err != nil {
    84  			return "", err
    85  		}
    86  		return randInt(f, t)
    87  	default:
    88  		return "", fmt.Errorf("maximum 2 arguments expected but got %d", len(args))
    89  	}
    90  }
    91  
    92  func randInt(f, t int64) (string, error) {
    93  	if t < f {
    94  		t, f = f, t
    95  	}
    96  	if f == 0 && t == 0 {
    97  		t = defaultMaxRandValue
    98  	}
    99  	if t == f {
   100  		f = t + defaultMaxRandValue
   101  	}
   102  	n := rand.Int63n(t - f)
   103  	n += f
   104  	return strconv.FormatInt(n, 10), nil
   105  }
   106  
   107  func RandString(args ...any) (string, error) {
   108  	switch len(args) {
   109  	case 0:
   110  		return randString(0, "")
   111  	case 1:
   112  		return randString(args[0], "")
   113  	case 2:
   114  		return randString(args[0], str.FormatString(args[1]))
   115  	default:
   116  		return "", fmt.Errorf("maximum 2 arguments expected but got %d", len(args))
   117  	}
   118  }
   119  
   120  func randString(cnt any, letters string) (string, error) {
   121  	n, err := numbers.ParseInt(cnt)
   122  	if err != nil {
   123  		return "", err
   124  	}
   125  	if n == 0 {
   126  		n = 1
   127  	}
   128  	return str.RandStringRunes(n, letters), nil
   129  }
   130  
   131  func UUID(args ...any) (string, error) {
   132  	v, err := uuid.NewV4()
   133  	if err != nil {
   134  		return "", err
   135  	}
   136  	return v.String(), nil
   137  }