diff --git a/opts.go b/opts.go index 5554b34..c4b080c 100644 --- a/opts.go +++ b/opts.go @@ -58,10 +58,6 @@ var ( func init() { for _, f := range parser.Functions { - // We skip variadic functions for now. - if f.Variadic != 0 { - continue - } if slices.Contains(f.ArgTypes, parser.ValueTypeString) { continue } diff --git a/walk.go b/walk.go index 79c5528..f98dc15 100644 --- a/walk.go +++ b/walk.go @@ -244,16 +244,20 @@ func (s *PromQLSmith) walkCall(valueTypes ...parser.ValueType) parser.Expr { } sort.Slice(funcs, func(i, j int) bool { return strings.Compare(funcs[i].Name, funcs[j].Name) < 0 }) expr.Func = funcs[s.rnd.Intn(len(funcs))] - s.walkFuncArgs(expr) + s.walkFunctions(expr) return expr } -func (s *PromQLSmith) walkFuncArgs(expr *parser.Call) { +func (s *PromQLSmith) walkFunctions(expr *parser.Call) { expr.Args = make([]parser.Expr, len(expr.Func.ArgTypes)) if expr.Func.Name == "holt_winters" { s.walkHoltWinters(expr) return } + if expr.Func.Variadic != 0 { + s.walkVariadicFunctions(expr) + return + } for i, arg := range expr.Func.ArgTypes { expr.Args[i] = s.Walk(arg) } @@ -265,6 +269,24 @@ func (s *PromQLSmith) walkHoltWinters(expr *parser.Call) { expr.Args[2] = &parser.NumberLiteral{Val: getNonZeroFloat64(s.rnd)} } +// Supported variadic functions include: +// days_in_month, day_of_month, day_of_week, day_of_year, year, +// hour, minute, month, round. +// Unsupported variadic functions include: +// label_join, sort_by_label_desc, sort_by_label +func (s *PromQLSmith) walkVariadicFunctions(expr *parser.Call) { + switch expr.Func.Name { + case "round": + expr.Args[0] = s.Walk(expr.Func.ArgTypes[0]) + expr.Args[1] = &parser.NumberLiteral{Val: float64(s.rnd.Intn(10))} + default: + // Rest of supported functions have either 0 or 1 function argument. + // If not specified it uses current timestamp instead of the vector timestamp. + // To reduce test flakiness we always use vector timestamp. + expr.Args[0] = s.Walk(expr.Func.ArgTypes[0]) + } +} + func (s *PromQLSmith) walkVectorSelector() parser.Expr { expr := &parser.VectorSelector{} expr.LabelMatchers = s.walkLabelMatchers() diff --git a/walk_test.go b/walk_test.go index 1fcec07..a5e0a91 100644 --- a/walk_test.go +++ b/walk_test.go @@ -396,17 +396,13 @@ func TestWalkSubQueryExpr(t *testing.T) { } } -func TestWalkFuncArgs(t *testing.T) { +func TestWalkFunctions(t *testing.T) { rnd := rand.New(rand.NewSource(time.Now().Unix())) opts := []Option{WithEnableOffset(true), WithEnableAtModifier(true)} p := New(rnd, testSeriesSet, opts...) - for _, f := range parser.Functions { - // Skip string type arg function for now as we don't support it. - if slices.Contains(f.ArgTypes, parser.ValueTypeString) { - continue - } + for _, f := range defaultSupportedFuncs { call := &parser.Call{Func: f} - p.walkFuncArgs(call) + p.walkFunctions(call) for i, arg := range call.Args { require.Equal(t, f.ArgTypes[i], arg.Type()) }