add support for variadic functions with no string type argument

Signed-off-by: Ben Ye <[email protected]>

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
 		}

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()

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())
 		}