Skip to content

Commit 4544b9e

Browse files
dvyukovdvyukov
committed
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.
1 parent d09a306 commit 4544b9e

File tree

9 files changed

+3338
-1355
lines changed

9 files changed

+3338
-1355
lines changed

‎pkg/aflow/flow_test.go‎

Lines changed: 245 additions & 1254 deletions
Large diffs are not rendered by default.

‎pkg/aflow/func_tool_test.go‎

Lines changed: 36 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,9 @@
44
package aflow
55

66
import (
7-
"context"
87
"errors"
9-
"path/filepath"
108
"testing"
11-
"time"
129

13-
"github.com/google/syzkaller/pkg/aflow/trajectory"
14-
"github.com/stretchr/testify/assert"
15-
"github.com/stretchr/testify/require"
1610
"google.golang.org/genai"
1711
)
1812

@@ -23,90 +17,43 @@ func TestToolErrors(t *testing.T) {
2317
type toolArgs struct {
2418
CallError bool `jsonschema:"call error"`
2519
}
26-
flows := make(map[string]*Flow)
27-
err := register[struct{}, flowOutputs]("test", "description", flows, []*Flow{
28-
{
29-
Root: &LLMAgent{
30-
Name: "smarty",
31-
Model: "model",
32-
Reply: "Reply",
33-
Temperature: 0,
34-
Instruction: "Do something!",
35-
Prompt: "Prompt",
36-
Tools: []Tool{
37-
NewFuncTool("faulty", func(ctx *Context, state struct{}, args toolArgs) (struct{}, error) {
38-
if args.CallError {
39-
return struct{}{}, BadCallError("you are wrong")
40-
}
41-
return struct{}{}, errors.New("hard error")
42-
}, "tool 1 description"),
43-
},
20+
testFlow[struct{}, flowOutputs](t, nil,
21+
"tool faulty failed: error: hard error\nargs: map[CallError:false]",
22+
&LLMAgent{
23+
Name: "smarty",
24+
Model: "model",
25+
Reply: "Reply",
26+
Temperature: 0,
27+
Instruction: "Do something!",
28+
Prompt: "Prompt",
29+
Tools: []Tool{
30+
NewFuncTool("faulty", func(ctx *Context, state struct{}, args toolArgs) (struct{}, error) {
31+
if args.CallError {
32+
return struct{}{}, BadCallError("you are wrong")
33+
}
34+
return struct{}{}, errors.New("hard error")
35+
}, "tool 1 description"),
4436
},
4537
},
46-
})
47-
require.NoError(t, err)
48-
replySeq := 0
49-
stub := &stubContext{
50-
// nolint:dupl
51-
generateContent: func(model string, cfg *genai.GenerateContentConfig, req []*genai.Content) (
52-
*genai.GenerateContentResponse, error) {
53-
replySeq++
54-
switch replySeq {
55-
case 1:
56-
return &genai.GenerateContentResponse{
57-
Candidates: []*genai.Candidate{{
58-
Content: &genai.Content{
59-
Role: string(genai.RoleModel),
60-
Parts: []*genai.Part{
61-
{
62-
FunctionCall: &genai.FunctionCall{
63-
ID: "id0",
64-
Name: "faulty",
65-
Args: map[string]any{
66-
"CallError": true,
67-
},
68-
},
69-
},
70-
}}}}}, nil
71-
case 2:
72-
assert.Equal(t, req[2], &genai.Content{
73-
Role: string(genai.RoleUser),
74-
Parts: []*genai.Part{
75-
{
76-
FunctionResponse: &genai.FunctionResponse{
77-
ID: "id0",
78-
Name: "faulty",
79-
Response: map[string]any{
80-
"error": "you are wrong",
81-
},
82-
},
83-
}}})
84-
return &genai.GenerateContentResponse{
85-
Candidates: []*genai.Candidate{{
86-
Content: &genai.Content{
87-
Role: string(genai.RoleModel),
88-
Parts: []*genai.Part{
89-
{
90-
FunctionCall: &genai.FunctionCall{
91-
ID: "id0",
92-
Name: "faulty",
93-
Args: map[string]any{
94-
"CallError": false,
95-
},
96-
},
97-
},
98-
}}}}}, nil
99-
default:
100-
t.Fatal("unexpected LLM calls")
101-
return nil, nil
102-
}
38+
[]any{
39+
&genai.Part{
40+
FunctionCall: &genai.FunctionCall{
41+
ID: "id0",
42+
Name: "faulty",
43+
Args: map[string]any{
44+
"CallError": true,
45+
},
46+
},
47+
},
48+
&genai.Part{
49+
FunctionCall: &genai.FunctionCall{
50+
ID: "id0",
51+
Name: "faulty",
52+
Args: map[string]any{
53+
"CallError": false,
54+
},
55+
},
56+
},
10357
},
104-
}
105-
ctx := context.WithValue(context.Background(), stubContextKey, stub)
106-
workdir := t.TempDir()
107-
cache, err := newTestCache(t, filepath.Join(workdir, "cache"), 0, time.Now)
108-
require.NoError(t, err)
109-
onEvent := func(span *trajectory.Span) error { return nil }
110-
_, err = flows["test"].Execute(ctx, "", workdir, nil, cache, onEvent)
111-
require.Equal(t, err.Error(), "tool faulty failed: error: hard error\nargs: map[CallError:false]")
58+
)
11259
}

‎pkg/aflow/testdata/TestToolErrors.llm.json‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
[
2+
{
3+
"Model": "model",
4+
"Config": {
5+
"systemInstruction": {
6+
"parts": [
7+
{
8+
"text": "Do something!\nPrefer calling several tools at the same time to save round-trips.\n"
9+
}
10+
],
11+
"role": "user"
12+
},
13+
"temperature": 0,
14+
"tools": [
15+
{
16+
"functionDeclarations": [
17+
{
18+
"description": "tool 1 description",
19+
"name": "faulty",
20+
"parametersJsonSchema": {
21+
"additionalProperties": false,
22+
"properties": {
23+
"CallError": {
24+
"description": "call error",
25+
"type": "boolean"
26+
}
27+
},
28+
"required": [
29+
"CallError"
30+
],
31+
"type": "object"
32+
},
33+
"responseJsonSchema": {
34+
"additionalProperties": false,
35+
"type": "object"
36+
}
37+
}
38+
]
39+
}
40+
],
41+
"responseModalities": [
42+
"TEXT"
43+
]
44+
},
45+
"Request": [
46+
{
47+
"parts": [
48+
{
49+
"text": "Prompt"
50+
}
51+
],
52+
"role": "user"
53+
}
54+
]
55+
},
56+
{
57+
"Model": "model",
58+
"Config": {
59+
"systemInstruction": {
60+
"parts": [
61+
{
62+
"text": "Do something!\nPrefer calling several tools at the same time to save round-trips.\n"
63+
}
64+
],
65+
"role": "user"
66+
},
67+
"temperature": 0,
68+
"tools": [
69+
{
70+
"functionDeclarations": [
71+
{
72+
"description": "tool 1 description",
73+
"name": "faulty",
74+
"parametersJsonSchema": {
75+
"additionalProperties": false,
76+
"properties": {
77+
"CallError": {
78+
"description": "call error",
79+
"type": "boolean"
80+
}
81+
},
82+
"required": [
83+
"CallError"
84+
],
85+
"type": "object"
86+
},
87+
"responseJsonSchema": {
88+
"additionalProperties": false,
89+
"type": "object"
90+
}
91+
}
92+
]
93+
}
94+
],
95+
"responseModalities": [
96+
"TEXT"
97+
]
98+
},
99+
"Request": [
100+
{
101+
"parts": [
102+
{
103+
"text": "Prompt"
104+
}
105+
],
106+
"role": "user"
107+
},
108+
{
109+
"parts": [
110+
{
111+
"functionCall": {
112+
"id": "id0",
113+
"args": {
114+
"CallError": true
115+
},
116+
"name": "faulty"
117+
}
118+
}
119+
],
120+
"role": "user"
121+
},
122+
{
123+
"parts": [
124+
{
125+
"functionResponse": {
126+
"id": "id0",
127+
"name": "faulty",
128+
"response": {
129+
"error": "you are wrong"
130+
}
131+
}
132+
}
133+
],
134+
"role": "user"
135+
}
136+
]
137+
}
138+
]

0 commit comments

Comments
 (0)