-
Notifications
You must be signed in to change notification settings - Fork 0
/
generator.go
87 lines (71 loc) · 2.37 KB
/
generator.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
package agent
import (
"context"
"errors"
"fmt"
"log"
"math/rand"
"strconv"
"time"
"github.com/mediocregopher/radix/v4"
"github.com/spf13/cobra"
)
var keyPrefix, prefixAmount []string
var entryCount int
var generatorCmd = &cobra.Command{
Use: "generate <redis>",
Short: "Create random entries in redis instance",
Long: `Create random entries in redis instance
Url can be provided as just "<host>:<port>" or in Redis URL format: "redis://[:<password>@]<host>:<port>[/<dbIndex>]"`,
Args: cobra.MinimumNArgs(1),
Run: func(cmd *cobra.Command, args []string) {
fmt.Println("Start generating keys")
ctx := context.Background()
randomMap, err := createRandomMap(keyPrefix, prefixAmount)
if err != nil {
log.Fatal(err)
}
fmt.Println("Generated random values: ", args[0])
generatorClient := createClient(ctx, sourcePassword, args[0], sourceUseTLS)
rand.Seed(time.Now().UTC().UnixNano())
for j := 0; j < entryCount; j++ {
for prefix, number := range randomMap {
for i := 0; i < number; i++ {
randVal := strconv.Itoa(rand.Int())
action := radix.Cmd(nil, "SET", prefix+randVal, randVal)
err = generatorClient.Do(ctx, action)
if err != nil {
fmt.Println(err)
}
}
}
fmt.Printf("Generation: %d done\n", j)
}
},
}
func createRandomMap(prefix []string, prefixAmount []string) (map[string]int, error) {
randomMap := make(map[string]int)
for key, val := range prefix {
randomMap[val] = 1
if key < len(prefixAmount) {
// parse to int because int array not possible via cli
countForPrefix, err := strconv.Atoi(prefixAmount[key])
if err != nil {
return nil, err
}
if countForPrefix <= 0 {
return nil, errors.New("count cannot be zero or negative")
}
randomMap[val] = countForPrefix
}
}
return randomMap, nil
}
func init() {
RootCmd.AddCommand(generatorCmd)
generatorCmd.Flags().BoolVar(&sourceUseTLS, "useTLS", true, "Enable TLS - default true")
generatorCmd.Flags().StringVar(&sourcePassword, "password", "", "Password for redis")
generatorCmd.Flags().StringArrayVar(&keyPrefix, "prefixes", []string{"mykey:", "testkey:"}, "List of prefixes for generated keys")
generatorCmd.Flags().StringArrayVar(&prefixAmount, "prefixAmount", []string{"1", "2"}, "Amount of keys to create for each prefix in one iteration")
generatorCmd.Flags().IntVar(&entryCount, "entryCount", 1, "Iteration count to perform")
}