Skip to content

Commit e6945f6

Browse files
BenjaminKazemiBenjaminKazemi
committed
chore(vertexai): Improve the test helpers.
1 parent 80a56cd commit e6945f6

File tree

7 files changed

+189
-197
lines changed

7 files changed

+189
-197
lines changed

‎vertexai/genai/agentengines_test.go‎

Lines changed: 23 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -15,89 +15,15 @@
1515
package genai
1616

1717
import (
18-
"flag"
1918
"fmt"
20-
"strings"
2119
"testing"
2220
"time"
2321

22+
"github.com/google/go-cmp/cmp"
23+
"github.com/google/go-cmp/cmp/cmpopts"
2424
"google.golang.org/genai"
2525
)
2626

27-
const (
28-
apiMode = "api" // API mode runs the tests in any environment where the tests can hit the actual service.
29-
unitMode = "unit" // Unit mode runs the test in the github actions using the mocked service (this is the default).
30-
)
31-
32-
var mode = flag.String("mode", unitMode, "Test mode")
33-
34-
func waitForAgentEngineOperation(tb testing.TB, name string, done bool, c *Client) any {
35-
tb.Helper()
36-
var res any
37-
for !done {
38-
tb.Logf("Waiting for operation to complete: [%s]\n", name)
39-
time.Sleep(5 * time.Second)
40-
op, err := c.AgentEngines.getAgentOperation(tb.Context(), name, nil)
41-
if err != nil {
42-
tb.Fatalf("getAgentOperation failed, err: %v", err)
43-
}
44-
done = op.Done
45-
res = op
46-
}
47-
return res
48-
}
49-
50-
func newTestClient(tb testing.TB) *Client {
51-
tb.Helper()
52-
client, err := NewGenAIClient(tb.Context(), &genai.ClientConfig{
53-
Backend: genai.BackendVertexAI,
54-
})
55-
if err != nil {
56-
tb.Fatal(err)
57-
}
58-
return client
59-
}
60-
61-
func getResourceNameFromOperation(operationName string) string {
62-
i := strings.Index(operationName, "/operations/")
63-
return operationName[:i]
64-
}
65-
66-
func createAgentEngineAndWait(t testing.TB, tt testing.TB, client *Client, config *CreateAgentEngineConfig) *ReasoningEngine {
67-
tt.Helper()
68-
if config == nil {
69-
config = &CreateAgentEngineConfig{
70-
DisplayName: tt.Name(),
71-
Description: "You can remove this agent engine if it is older than 10 minutes. It must be an orphan AE.",
72-
}
73-
}
74-
if config.DisplayName == "" {
75-
config.DisplayName = tt.Name()
76-
}
77-
if config.Description == "" {
78-
config.Description = "You can remove this agent engine if it is older than 10 minutes. It must be an orphan AE."
79-
}
80-
createOp, err := client.AgentEngines.create(tt.Context(), config)
81-
if err != nil {
82-
tt.Fatalf("create() failed unexpectedly: %v", err)
83-
}
84-
tt.Cleanup(cleanupAgentEngine(t, client, getResourceNameFromOperation(createOp.Name)))
85-
createOp = waitForAgentEngineOperation(tt, createOp.Name, createOp.Done, client).(*AgentEngineOperation)
86-
return createOp.Response
87-
}
88-
89-
func cleanupAgentEngine(t testing.TB, client *Client, name string) func() {
90-
return func() {
91-
t.Logf("Cleaning up AgentEngine: %s", name)
92-
deleteOp, err := client.AgentEngines.delete(t.Context(), name, genai.Ptr(true), nil)
93-
if err != nil {
94-
t.Logf("cleanup() failed, err: %v", err)
95-
} else {
96-
waitForAgentEngineOperation(t, deleteOp.Name, deleteOp.Done, client)
97-
}
98-
}
99-
}
100-
10127
func TestAgentEngines(t *testing.T) {
10228
if *mode != apiMode {
10329
t.Skipf("Skipping %s. We only tun these in the api mode.", t.Name())
@@ -163,13 +89,14 @@ func TestAgentEngines(t *testing.T) {
16389
ctx := tt.Context()
16490
client := newTestClient(t)
16591
re := createAgentEngineAndWait(t, tt, client, nil)
166-
16792
deleteOp, err := client.AgentEngines.delete(t.Context(), re.Name, nil, nil)
16893
if err != nil {
16994
tt.Fatalf("delete() failed unexpectedly: %v", err)
17095
}
171-
waitForAgentEngineOperation(t, deleteOp.Name, deleteOp.Done, client)
172-
96+
operation := func() (*AgentEngineOperation, error) {
97+
return client.AgentEngines.getAgentOperation(tt.Context(), deleteOp.Name, nil)
98+
}
99+
waitForOperation(t, operation)
173100
got, err := client.AgentEngines.get(ctx, re.Name, nil)
174101
if err == nil {
175102
t.Errorf("delete() didn't remove the reasoning engine, want error(NOT_FOUND), got: %v", got)
@@ -179,11 +106,23 @@ func TestAgentEngines(t *testing.T) {
179106
t.Run("Get", func(tt *testing.T) {
180107
ctx := tt.Context()
181108
client := newTestClient(tt)
182-
re := createAgentEngineAndWait(t, tt, client, nil)
183-
_, err := client.AgentEngines.get(ctx, re.Name, nil)
109+
want := &ReasoningEngine{
110+
DisplayName: tt.Name(),
111+
Description: "You can remove this agent engine if it is older than 10 minutes. It must be an orphan AE.",
112+
}
113+
config := &CreateAgentEngineConfig{
114+
DisplayName: want.DisplayName,
115+
Description: want.Description,
116+
}
117+
re := createAgentEngineAndWait(t, tt, client, config)
118+
got, err := client.AgentEngines.get(ctx, re.Name, nil)
184119
if err != nil {
185120
tt.Errorf("get() failed unexpectedly: %v", err)
186121
}
122+
if diff := cmp.Diff(got, want, cmpopts.IgnoreFields(ReasoningEngine{}, "CreateTime", "Spec", "UpdateTime", "Name")); diff != "" {
123+
tt.Errorf("create() and get() had diff (-got +want): %v", diff)
124+
}
125+
187126
})
188127

189128
t.Run("List", func(tt *testing.T) {
@@ -210,11 +149,10 @@ func TestAgentEngines(t *testing.T) {
210149
if err != nil {
211150
tt.Fatalf("update() failed unexpectedly: %v", err)
212151
}
213-
waitForAgentEngineOperation(tt, op.Name, op.Done, client)
214-
updated, err := client.AgentEngines.get(ctx, re.Name, nil)
215-
if err != nil {
216-
tt.Errorf("get() failed unexpectedly: %v", err)
152+
operation := func() (*AgentEngineOperation, error) {
153+
return client.AgentEngines.getAgentOperation(tt.Context(), op.Name, nil)
217154
}
155+
updated := waitForOperation(tt, operation).Response
218156
if got := updated.DisplayName; got != want {
219157
tt.Errorf("update() returned DisplayName %v, want %v", got, want)
220158
}

‎vertexai/genai/memories_test.go‎

Lines changed: 1 addition & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -23,39 +23,6 @@ import (
2323
"google.golang.org/genai"
2424
)
2525

26-
func waitForMemoryOperation(tb testing.TB, name string, done bool, c *Client) any {
27-
tb.Helper()
28-
var res any
29-
for !done {
30-
tb.Logf("Waiting for operation to complete: [%s]\n", name)
31-
time.Sleep(5 * time.Second)
32-
op, err := c.AgentEngines.Memories.getMemoryOperation(tb.Context(), name, nil)
33-
if err != nil {
34-
tb.Fatalf("getAgentOperation failed, err: %v", err)
35-
}
36-
done = op.Done
37-
res = op
38-
}
39-
return res
40-
}
41-
42-
func createAgentEngineMemoryAndWait(tt testing.TB, client *Client, name string, m *Memory) *Memory {
43-
tt.Helper()
44-
config := &AgentEngineMemoryConfig{
45-
DisplayName: m.DisplayName,
46-
Description: m.Description,
47-
Metadata: m.Metadata,
48-
}
49-
createOp, err := client.AgentEngines.Memories.create(tt.Context(), name, m.Fact, m.Scope, config)
50-
if err != nil {
51-
tt.Fatalf("create() failed unexpectedly: %v", err)
52-
}
53-
if !createOp.Done {
54-
waitForMemoryOperation(tt, createOp.Name, createOp.Done, client)
55-
}
56-
return createOp.Response
57-
}
58-
5926
func TestAgentEngineMemories(t *testing.T) {
6027
if *mode != apiMode {
6128
t.Skipf("Skipping %s. We only tun these in the api mode.", t.Name())
@@ -80,7 +47,7 @@ func TestAgentEngineMemories(t *testing.T) {
8047
"my_timestamp_key": {TimestampValue: timestamp},
8148
},
8249
}
83-
response := createAgentEngineMemoryAndWait(tt, client, re.Name, want)
50+
response := createAgentEngineMemoryAndWait(tt, client, re, want)
8451
got, err := client.AgentEngines.Memories.Get(tt.Context(), response.Name, nil)
8552
if err != nil {
8653
tt.Fatalf("get() failed unexpectedly: %v", err)

‎vertexai/genai/memoryrevisions_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ func TestAgentEngineMemoryRevisions(t *testing.T) {
4545
"my_timestamp_key": {TimestampValue: timestamp},
4646
},
4747
}
48-
response := createAgentEngineMemoryAndWait(tt, client, re.Name, want)
48+
response := createAgentEngineMemoryAndWait(tt, client, re, want)
4949
got, err := client.AgentEngines.Memories.Revisions.list(tt.Context(), response.Name, nil)
5050
if err != nil {
5151
tt.Fatalf("get() failed unexpectedly: %v", err)

‎vertexai/genai/sandboxes_test.go‎

Lines changed: 1 addition & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -22,34 +22,6 @@ import (
2222
"github.com/google/go-cmp/cmp/cmpopts"
2323
)
2424

25-
func waitForSandboxesOperation(tt testing.TB, name string, done bool, c *Client) *AgentEngineSandboxOperation {
26-
tt.Helper()
27-
var res *AgentEngineSandboxOperation
28-
for !done {
29-
tt.Logf("Waiting for operation to complete: [%s]\n", name)
30-
time.Sleep(5 * time.Second)
31-
op, err := c.AgentEngines.Sandboxes.getSandboxOperation(tt.Context(), name, nil)
32-
if err != nil {
33-
tt.Fatalf("getAgentOperation failed, err: %v", err)
34-
}
35-
done = op.Done
36-
res = op
37-
}
38-
return res
39-
}
40-
41-
func createAgentEngineSandboxesAndWait(tt testing.TB, client *Client, name string, spec *SandboxEnvironmentSpec, config *CreateAgentEngineSandboxConfig) *SandboxEnvironment {
42-
tt.Helper()
43-
createOp, err := client.AgentEngines.Sandboxes.create(tt.Context(), name, spec, config)
44-
if err != nil {
45-
tt.Fatalf("create() failed unexpectedly: %v", err)
46-
}
47-
if !createOp.Done {
48-
createOp = waitForSandboxesOperation(tt, createOp.Name, createOp.Done, client)
49-
}
50-
return createOp.Response
51-
}
52-
5325
func TestAgentEngineSandboxes(t *testing.T) {
5426
if *mode != apiMode {
5527
t.Skipf("Skipping %s. We only tun these in the api mode.", t.Name())
@@ -72,7 +44,7 @@ func TestAgentEngineSandboxes(t *testing.T) {
7244
DisplayName: want.DisplayName,
7345
TTL: want.TTL,
7446
}
75-
got := createAgentEngineSandboxesAndWait(tt, client, re.Name, want.Spec, config)
47+
got := createAgentEngineSandboxesAndWait(tt, client, re, want.Spec, config)
7648
if diff := cmp.Diff(got, want, cmpopts.IgnoreFields(SandboxEnvironment{}, "Name", "ExpireTime", "TTL")); diff != "" {
7749
tt.Errorf("create() had diff (-got +want): %v", diff)
7850
}

‎vertexai/genai/sessionevents_test.go‎

Lines changed: 2 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -23,17 +23,6 @@ import (
2323
"google.golang.org/genai"
2424
)
2525

26-
func createAgentEngineSessionEventAndWait(tt testing.TB, client *Client, name string, se *SessionEvent) {
27-
tt.Helper()
28-
config := &AppendAgentEngineSessionEventConfig{
29-
Content: se.Content,
30-
}
31-
_, err := client.AgentEngines.Sessions.Events.Append(tt.Context(), name, se.Author, se.InvocationID, se.Timestamp, config)
32-
if err != nil {
33-
tt.Fatalf("append() failed unexpectedly: %v", err)
34-
}
35-
}
36-
3726
func TestAgentEngineSessionEvents(t *testing.T) {
3827
if *mode != apiMode {
3928
t.Skipf("Skipping %s. We only tun these in the api mode.", t.Name())
@@ -50,7 +39,7 @@ func TestAgentEngineSessionEvents(t *testing.T) {
5039
TTL: 24 * time.Hour,
5140
UserID: "test-user-123",
5241
}
53-
session = createAgentEngineSessionAndWait(tt, client, re.Name, session)
42+
session = createAgentEngineSessionAndWait(tt, client, re, session)
5443
timestamp, err := time.Parse(time.RFC3339Nano, "2026-03-10T15:30:45.0Z")
5544
if err != nil {
5645
tt.Fatalf("Error parsing time, err: %v", err)
@@ -66,7 +55,7 @@ func TestAgentEngineSessionEvents(t *testing.T) {
6655
}},
6756
},
6857
}
69-
createAgentEngineSessionEventAndWait(tt, client, session.Name, want)
58+
createAgentEngineSessionEvent(tt, client, session, want)
7059

7160
got, err := client.AgentEngines.Sessions.Events.list(tt.Context(), session.Name, nil)
7261
if err != nil {

‎vertexai/genai/sessions_test.go‎

Lines changed: 1 addition & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -22,40 +22,6 @@ import (
2222
"github.com/google/go-cmp/cmp/cmpopts"
2323
)
2424

25-
func waitForSessionsOperation(tt testing.TB, name string, done bool, c *Client) any {
26-
tt.Helper()
27-
var res any
28-
for !done {
29-
tt.Logf("Waiting for operation to complete: [%s]\n", name)
30-
time.Sleep(5 * time.Second)
31-
op, err := c.AgentEngines.Sessions.getSessionOperation(tt.Context(), name, nil)
32-
if err != nil {
33-
tt.Fatalf("getAgentOperation failed, err: %v", err)
34-
}
35-
done = op.Done
36-
res = op
37-
}
38-
return res
39-
}
40-
41-
func createAgentEngineSessionAndWait(tt testing.TB, client *Client, name string, s *Session) *Session {
42-
tt.Helper()
43-
config := &CreateAgentEngineSessionConfig{
44-
DisplayName: s.DisplayName,
45-
SessionState: s.SessionState,
46-
TTL: s.TTL,
47-
Labels: s.Labels,
48-
}
49-
createOp, err := client.AgentEngines.Sessions.create(tt.Context(), name, s.UserID, config)
50-
if err != nil {
51-
tt.Fatalf("create() failed unexpectedly: %v", err)
52-
}
53-
if !createOp.Done {
54-
createOp = waitForSessionsOperation(tt, createOp.Name, createOp.Done, client).(*AgentEngineSessionOperation)
55-
}
56-
return createOp.Response
57-
}
58-
5925
func TestAgentEngineSessions(t *testing.T) {
6026
if *mode != apiMode {
6127
t.Skipf("Skipping %s. We only tun these in the api mode.", t.Name())
@@ -71,7 +37,7 @@ func TestAgentEngineSessions(t *testing.T) {
7137
TTL: 24 * time.Hour,
7238
UserID: "test-user-123",
7339
}
74-
got := createAgentEngineSessionAndWait(tt, client, re.Name, want)
40+
got := createAgentEngineSessionAndWait(tt, client, re, want)
7541
if diff := cmp.Diff(got, want, cmpopts.IgnoreFields(Session{}, "Name")); diff != "" {
7642
tt.Errorf("create() had diff (-got +want): %v", diff)
7743
}

0 commit comments

Comments
 (0)