1package provider
2
3import (
4 "net/http"
5 "os"
6 "path/filepath"
7 "testing"
8)
9
10func TestVertexConfigParsesAuthorizedUser(t *testing.T) {
11 // Write a fake ADC user-OAuth file and point GOOGLE_APPLICATION_CREDENTIALS at it.
12 tmp := t.TempDir()
13 path := filepath.Join(tmp, "adc.json")
14 body := `{
15 "type": "authorized_user",
16 "client_id": "fake.apps.googleusercontent.com",
17 "client_secret": "fake-secret",
18 "refresh_token": "1//fake-refresh"
19 }`
20 if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
21 t.Fatal(err)
22 }
23 t.Setenv("GOOGLE_APPLICATION_CREDENTIALS", path)
24 t.Setenv("GOOGLE_CLOUD_PROJECT", "test-proj")
25 t.Setenv("GOOGLE_CLOUD_API_KEY", "") // ensure SA / user path is chosen
26 cfg, err := loadVertexConfig()
27 if err != nil {
28 t.Fatalf("loadVertexConfig: %v", err)
29 }
30 if cfg.userClientID != "fake.apps.googleusercontent.com" {
31 t.Errorf("client_id = %q", cfg.userClientID)
32 }
33 if cfg.userRefreshToken != "1//fake-refresh" {
34 t.Errorf("refresh_token = %q", cfg.userRefreshToken)
35 }
36 if cfg.userTokenURI != "https://oauth2.googleapis.com/token" {
37 t.Errorf("token_uri default not applied: %q", cfg.userTokenURI)
38 }
39 if cfg.cacheKey() != "user:fake.apps.googleusercontent.com" {
40 t.Errorf("cacheKey = %q", cfg.cacheKey())
41 }
42}
43
44type vertexRoundTripFunc func(*http.Request) (*http.Response, error)
45
46func (f vertexRoundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
47 return f(req)
48}
49
50func TestVertexTransportEndpoint(t *testing.T) {
51 tests := []struct {
52 name string
53 location string
54 wantHost string
55 }{
56 {name: "global", location: "global", wantHost: "aiplatform.googleapis.com"},
57 {name: "regional", location: "us-central1", wantHost: "us-central1-aiplatform.googleapis.com"},
58 }
59
60 for _, tt := range tests {
61 t.Run(tt.name, func(t *testing.T) {
62 var got *http.Request
63 transport := &vertexTransport{
64 cfg: &vertexConfig{
65 project: "test-project",
66 location: tt.location,
67 apiKey: "test-key",
68 },
69 inner: vertexRoundTripFunc(func(req *http.Request) (*http.Response, error) {
70 got = req
71 return &http.Response{StatusCode: http.StatusOK}, nil
72 }),
73 }
74 req, err := http.NewRequest(http.MethodPost, "https://placeholder.example/v1beta/models/gemini-test:streamGenerateContent?alt=sse", nil)
75 if err != nil {
76 t.Fatal(err)
77 }
78
79 if _, err := transport.RoundTrip(req); err != nil {
80 t.Fatalf("RoundTrip: %v", err)
81 }
82 if got.URL.Host != tt.wantHost {
83 t.Errorf("host = %q, want %q", got.URL.Host, tt.wantHost)
84 }
85 wantPath := "/v1/projects/test-project/locations/" + tt.location + "/publishers/google/models/gemini-test:streamGenerateContent"
86 if got.URL.Path != wantPath {
87 t.Errorf("path = %q, want %q", got.URL.Path, wantPath)
88 }
89 })
90 }
91}
92
93func TestVertexConfigRejectsBadType(t *testing.T) {
94 tmp := t.TempDir()
95 path := filepath.Join(tmp, "adc.json")
96 body := `{"type": "something_else"}`
97 if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
98 t.Fatal(err)
99 }
100 t.Setenv("GOOGLE_APPLICATION_CREDENTIALS", path)
101 t.Setenv("GOOGLE_CLOUD_PROJECT", "test-proj")
102 t.Setenv("GOOGLE_CLOUD_API_KEY", "")
103 _, err := loadVertexConfig()
104 if err == nil {
105 t.Fatal("expected error for unknown credential type")
106 }
107}