pkg/aflow: refactor tests

Add helper function that executes test workflows,
compares results (trajectory, LLM requests) against golden files,
and if requested updates these golden files.

Author: dvyukov

pkg/aflow/flow_test.go

@@ -4,15 +4,9 @@
 package aflow
 
 import (
-	"context"
 	"errors"
-	"path/filepath"
 	"testing"
-	"time"
 
-	"github.com/google/syzkaller/pkg/aflow/trajectory"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"google.golang.org/genai"
 )
 
@@ -23,90 +17,43 @@ func TestToolErrors(t *testing.T) {
 	type toolArgs struct {
 		CallError bool `jsonschema:"call error"`
 	}
-	flows := make(map[string]*Flow)
-	err := register[struct{}, flowOutputs]("test", "description", flows, []*Flow{
-		{
-			Root: &LLMAgent{
-				Name:        "smarty",
-				Model:       "model",
-				Reply:       "Reply",
-				Temperature: 0,
-				Instruction: "Do something!",
-				Prompt:      "Prompt",
-				Tools: []Tool{
-					NewFuncTool("faulty", func(ctx *Context, state struct{}, args toolArgs) (struct{}, error) {
-						if args.CallError {
-							return struct{}{}, BadCallError("you are wrong")
-						}
-						return struct{}{}, errors.New("hard error")
-					}, "tool 1 description"),
-				},
+	testFlow[struct{}, flowOutputs](t, nil,
+		"tool faulty failed: error: hard error\nargs: map[CallError:false]",
+		&LLMAgent{
+			Name:        "smarty",
+			Model:       "model",
+			Reply:       "Reply",
+			Temperature: 0,
+			Instruction: "Do something!",
+			Prompt:      "Prompt",
+			Tools: []Tool{
+				NewFuncTool("faulty", func(ctx *Context, state struct{}, args toolArgs) (struct{}, error) {
+					if args.CallError {
+						return struct{}{}, BadCallError("you are wrong")
+					}
+					return struct{}{}, errors.New("hard error")
+				}, "tool 1 description"),
 			},
 		},
-	})
-	require.NoError(t, err)
-	replySeq := 0
-	stub := &stubContext{
-		// nolint:dupl
-		generateContent: func(model string, cfg *genai.GenerateContentConfig, req []*genai.Content) (
-			*genai.GenerateContentResponse, error) {
-			replySeq++
-			switch replySeq {
-			case 1:
-				return &genai.GenerateContentResponse{
-					Candidates: []*genai.Candidate{{
-						Content: &genai.Content{
-							Role: string(genai.RoleModel),
-							Parts: []*genai.Part{
-								{
-									FunctionCall: &genai.FunctionCall{
-										ID:   "id0",
-										Name: "faulty",
-										Args: map[string]any{
-											"CallError": true,
-										},
-									},
-								},
-							}}}}}, nil
-			case 2:
-				assert.Equal(t, req[2], &genai.Content{
-					Role: string(genai.RoleUser),
-					Parts: []*genai.Part{
-						{
-							FunctionResponse: &genai.FunctionResponse{
-								ID:   "id0",
-								Name: "faulty",
-								Response: map[string]any{
-									"error": "you are wrong",
-								},
-							},
-						}}})
-				return &genai.GenerateContentResponse{
-					Candidates: []*genai.Candidate{{
-						Content: &genai.Content{
-							Role: string(genai.RoleModel),
-							Parts: []*genai.Part{
-								{
-									FunctionCall: &genai.FunctionCall{
-										ID:   "id0",
-										Name: "faulty",
-										Args: map[string]any{
-											"CallError": false,
-										},
-									},
-								},
-							}}}}}, nil
-			default:
-				t.Fatal("unexpected LLM calls")
-				return nil, nil
-			}
+		[]any{
+			&genai.Part{
+				FunctionCall: &genai.FunctionCall{
+					ID:   "id0",
+					Name: "faulty",
+					Args: map[string]any{
+						"CallError": true,
+					},
+				},
+			},
+			&genai.Part{
+				FunctionCall: &genai.FunctionCall{
+					ID:   "id0",
+					Name: "faulty",
+					Args: map[string]any{
+						"CallError": false,
+					},
+				},
+			},
 		},
-	}
-	ctx := context.WithValue(context.Background(), stubContextKey, stub)
-	workdir := t.TempDir()
-	cache, err := newTestCache(t, filepath.Join(workdir, "cache"), 0, time.Now)
-	require.NoError(t, err)
-	onEvent := func(span *trajectory.Span) error { return nil }
-	_, err = flows["test"].Execute(ctx, "", workdir, nil, cache, onEvent)
-	require.Equal(t, err.Error(), "tool faulty failed: error: hard error\nargs: map[CallError:false]")
+	)
 }

pkg/aflow/func_tool_test.go

@@ -0,0 +1,118 @@
+// Copyright 2026 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package aflow
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/google/syzkaller/pkg/aflow/trajectory"
+	"github.com/google/syzkaller/pkg/osutil"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/genai"
+)
+
+var flagUpdate = flag.Bool("update", false, "update golden test files to match the actual execution")
+
+// testFlow executes the provided test workflow by returning LLM replies from llmReplies.
+// The result can be either a map[string]any with Outputs fields, or an error,
+// if an error is expected as the result of the execution.
+// llmReplies objects can be either *genai.Part, []*genai.Part, or an error.
+// Requests sent to LLM are compared against "testdata/TestName.llm.json" file.
+// Resulting trajectory is compared against "testdata/TestName.trajectory.json" file.
+// If -update flag is provided, the golden testdata files are updated to match the actual execution.
+func testFlow[Inputs, Outputs any](t *testing.T, inputs map[string]any, result any, root Action, llmReplies []any) {
+	flows := make(map[string]*Flow)
+	err := register[Inputs, Outputs]("test", "description", flows, []*Flow{{Root: root}})
+	require.NoError(t, err)
+	type llmRequest struct {
+		Model   string
+		Config  *genai.GenerateContentConfig
+		Request []*genai.Content
+	}
+	var requests []llmRequest
+	var stubTime time.Time
+	stub := &stubContext{
+		timeNow: func() time.Time {
+			stubTime = stubTime.Add(time.Second)
+			return stubTime
+		},
+		generateContent: func(model string, cfg *genai.GenerateContentConfig, req []*genai.Content) (
+			*genai.GenerateContentResponse, error) {
+			requests = append(requests, llmRequest{model, cfg, req})
+			require.NotEmpty(t, llmReplies, "unexpected LLM call")
+			reply := llmReplies[0]
+			llmReplies = llmReplies[1:]
+			switch reply := reply.(type) {
+			case error:
+				return nil, reply
+			case *genai.Part:
+				return &genai.GenerateContentResponse{
+					Candidates: []*genai.Candidate{{Content: &genai.Content{
+						Role:  string(genai.RoleUser),
+						Parts: []*genai.Part{reply},
+					}}}}, nil
+			case []*genai.Part:
+				return &genai.GenerateContentResponse{
+					Candidates: []*genai.Candidate{{Content: &genai.Content{
+						Role:  string(genai.RoleUser),
+						Parts: reply,
+					}}}}, nil
+			default:
+				t.Fatalf("bad LLM reply type %T", reply)
+				return nil, nil
+			}
+		},
+	}
+	var spans []*trajectory.Span
+	onEvent := func(span *trajectory.Span) error {
+		spans = append(spans, span)
+		return nil
+	}
+	ctx := context.WithValue(context.Background(), stubContextKey, stub)
+	workdir := t.TempDir()
+	cache, err := newTestCache(t, filepath.Join(workdir, "cache"), 0, time.Now)
+	require.NoError(t, err)
+	if inputs == nil {
+		inputs = map[string]any{}
+	}
+	got, err := flows["test"].Execute(ctx, "", workdir, inputs, cache, onEvent)
+	switch result := result.(type) {
+	case map[string]any:
+		require.NoError(t, err)
+		require.Equal(t, got, result)
+	case string:
+		require.Error(t, err)
+		require.Equal(t, err.Error(), result)
+	default:
+		t.Fatalf("bad result type %T", result)
+	}
+	// We need to pass spans/requests via double marshal/unmarshal round-trip
+	// b/c some values change during the first round-trip (int->float64, jsonschema).
+	spansData, err := json.Marshal(spans)
+	require.NoError(t, err)
+	spans = nil
+	require.NoError(t, json.Unmarshal(spansData, &spans))
+	requestsData, err := json.Marshal(requests)
+	require.NoError(t, err)
+	requests = nil
+	require.NoError(t, json.Unmarshal(requestsData, &requests))
+	trajectoryFile := filepath.Join("testdata", t.Name()+".trajectory.json")
+	requestsFile := filepath.Join("testdata", t.Name()+".llm.json")
+	if *flagUpdate {
+		require.NoError(t, osutil.WriteJSON(trajectoryFile, spans))
+		require.NoError(t, osutil.WriteJSON(requestsFile, requests))
+	}
+	wantSpans, err := osutil.ReadJSON[[]*trajectory.Span](trajectoryFile)
+	require.NoError(t, err)
+	require.Equal(t, spans, wantSpans)
+	wantRequests, err := osutil.ReadJSON[[]llmRequest](requestsFile)
+	require.NoError(t, err)
+	require.Equal(t, requests, wantRequests)
+	require.Empty(t, llmReplies)
+}