genai: make bucket-referencing test more resilient (#10571)

eliben · web-flow · commit ba8294215ff2 · 2024-07-23T11:16:40.000-07:00
diff --git a/vertexai/genai/caching_test.go b/vertexai/genai/caching_test.go
@@ -16,6 +16,7 @@ package genai
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"strings"
 	"testing"
@@ -76,8 +77,10 @@ func testCaching(t *testing.T, client *Client) {
 			return cc
 		}
 
+		vertexProject := os.Getenv("VERTEX_PROJECT_ID")
+
 		want := &CachedContent{
-			Model: "projects/" + os.Getenv("VERTEX_PROJECT_ID") +
+			Model: "projects/" + vertexProject +
 				"/locations/us-central1/publishers/google/models/" + model,
 			CreateTime: time.Now().UTC(),
 			UpdateTime: time.Now().UTC(),
@@ -93,13 +96,15 @@ func testCaching(t *testing.T, client *Client) {
 			}
 		}
 
+		gcsFilePath := fmt.Sprintf("gs://0002-test-multimodal-%s/embeddings/xray-embedding.json", vertexProject)
+
 		ttl := 30 * time.Minute
 		wantExpireTime := time.Now().Add(ttl)
 		argcc := &CachedContent{
 			Model:      model,
 			Expiration: ExpireTimeOrTTL{TTL: ttl},
 			Contents: []*Content{{Role: "user", Parts: []Part{
-				FileData{MIMEType: "text/plain", FileURI: "gs://0002-test-multimodal/embeddings/xray-embedding.json"},
+				FileData{MIMEType: "text/plain", FileURI: gcsFilePath},
 			}}},
 		}
 		cc := must(client.CreateCachedContent(ctx, argcc))
diff --git a/vertexai/genai/client_test.go b/vertexai/genai/client_test.go
@@ -678,3 +678,12 @@ func TestInferLocation(t *testing.T) {
 		})
 	}
 }
+
+func printResponse(resp *GenerateContentResponse) {
+	for _, cand := range resp.Candidates {
+		for _, part := range cand.Content.Parts {
+			fmt.Println(part)
+		}
+	}
+	fmt.Println("---")
+}

Original file line number	Diff line number	Diff line change
`@@ -678,3 +678,12 @@ func TestInferLocation(t *testing.T) {`
`678`	`678`	`})`
`679`	`679`	`}`
`680`	`680`	`}`
	`681`	`+`
	`682`	`+func printResponse(resp *GenerateContentResponse) {`
	`683`	`+ for _, cand := range resp.Candidates {`
	`684`	`+ for _, part := range cand.Content.Parts {`
	`685`	`+ fmt.Println(part)`
	`686`	`+ }`
	`687`	`+ }`
	`688`	`+ fmt.Println("---")`
	`689`	`+}`