/
generator.go
175 lines (153 loc) · 4.56 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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
package schema
import (
"errors"
"math/rand"
"strings"
"time"
"github.com/go-faker/faker/v4"
"github.com/go-faker/faker/v4/pkg/options"
"github.com/kuzxnia/loadbot/lbot/config"
"github.com/samber/lo"
"go.mongodb.org/mongo-driver/bson"
)
var (
DefaultGeneratorFieldMapper = NewGeneratorFieldMapper()
// todo: better validation
GeneratorFieldTypes = lo.Keys(DefaultGeneratorFieldMapper.FieldTypeMapper)
)
// todo: add interface
type GeneratorFieldMapper struct {
// todo: make private
FieldTypeMapper map[string]func(opts ...options.OptionFunc) string
}
func NewGeneratorFieldMapper() *GeneratorFieldMapper {
return &GeneratorFieldMapper{
FieldTypeMapper: map[string]func(opts ...options.OptionFunc) string{
"#id": faker.UUIDDigit,
"#string": faker.Word,
"#word": faker.Word,
// internet
"#email": faker.Email,
"#username": faker.Username,
"#password": faker.Password,
// person
"#name": faker.Name,
"#first_name": faker.FirstName,
"#first_name_male": faker.FirstNameMale,
"#first_name_female": faker.FirstNameFemale,
"#last_name": faker.LastName,
"#title_male": faker.TitleMale,
"#title_female": faker.TitleFemale,
"#phone_number": faker.Phonenumber,
},
}
}
func (m *GeneratorFieldMapper) Generate(field string) (result interface{}, err error) {
if generate, ok := m.FieldTypeMapper[field]; ok {
return generate(), nil
} else {
return nil, errors.New("Invalid field mapper, got: " + field)
}
}
func (m *GeneratorFieldMapper) Set(field string, valueGenerator func(opts ...options.OptionFunc) string) {
m.FieldTypeMapper[field] = valueGenerator
}
type DataGenerator interface {
Generate() (interface{}, error)
GenerateFromTemplate(interface{}) (interface{}, error)
}
func NewDataGenerator(schema *config.Schema, dataSize int) DataGenerator {
// todo: check size of object using, unsafe.Sizeof( )
if schema != nil {
return DataGenerator(
&StructuralizableDataGenerator{
schema: schema,
// add support for custom byte size
},
)
}
return DataGenerator(
&MeasurableDataGenerator{
dataSize: dataSize,
},
)
}
type MeasurableDataGenerator struct {
dataSize int
}
func (g *MeasurableDataGenerator) Generate() (interface{}, error) {
// size - 33
return &bson.M{"data": randStringBytes(g.dataSize)}, nil
}
func (g *MeasurableDataGenerator) GenerateFromTemplate(template interface{}) (interface{}, error) {
switch value := template.(type) {
case string:
return randStringBytes(g.dataSize), nil
case map[string]interface{}:
result := make(map[string]interface{})
for k, nestedTemplate := range value {
value, err := g.GenerateFromTemplate(nestedTemplate)
if err != nil {
return nil, err
}
result[k] = value
}
return result, nil
default:
return nil, errors.New("Invalid schema format")
}
}
const (
letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
letterIdxBits = 6 // 6 bits to represent a letter index
letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
letterIdxMax = 63 / letterIdxBits // # of letter indices fitting in 63 bits
)
var src = rand.NewSource(time.Now().UnixNano())
func randStringBytes(n int) string {
sb := strings.Builder{}
sb.Grow(n)
// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
for i, cache, remain := n-1, src.Int63(), letterIdxMax; i >= 0; {
if remain == 0 {
cache, remain = src.Int63(), letterIdxMax
}
if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
sb.WriteByte(letterBytes[idx])
i--
}
cache >>= letterIdxBits
remain--
}
return sb.String()
}
type StructuralizableDataGenerator struct {
schema *config.Schema
}
func (g *StructuralizableDataGenerator) Generate() (interface{}, error) {
result, error := g.GenerateFromTemplate(g.schema.Schema)
return result, error
}
// recurent func for parsing with building nested bson
func (g *StructuralizableDataGenerator) GenerateFromTemplate(template interface{}) (interface{}, error) {
switch value := template.(type) {
case string:
generatedValue, err := DefaultGeneratorFieldMapper.Generate(value)
if err != nil {
return nil, errors.New("Invalid field mapper, got: " + value)
}
return generatedValue, nil
case map[string]interface{}:
result := make(map[string]interface{})
for k, nestedTemplate := range value {
value, err := g.GenerateFromTemplate(nestedTemplate)
if err != nil {
return nil, err
}
result[k] = value
}
return result, nil
default:
return nil, errors.New("Invalid schema format")
}
}