Skip to content

Commit b1f949a

Browse files
BenjaminKazemicopybara-github
authored andcommitted
chore(vertexai): Add some tests
FUTURE_COPYBARA_INTEGRATE_REVIEW=#14351 from renovate-bot:renovate/main-deps 2e7e530 PiperOrigin-RevId: 895405821
1 parent 4d1b44e commit b1f949a

2 files changed

Lines changed: 48 additions & 0 deletions

File tree

vertexai/genai/client.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,8 +459,18 @@ func float32pToInt32p(x *float32) *int32 {
459459
return &i
460460
}
461461

462+
func validateClientConfig(cc *genai.ClientConfig) error {
463+
if cc != nil && cc.Backend == genai.BackendGeminiAPI {
464+
return fmt.Errorf("only Vertex AI backend is supported")
465+
}
466+
return nil
467+
}
468+
462469
// NewGenAIClient creates a new Google Vertex AI client and configures the the GenAI components.
463470
func NewGenAIClient(ctx context.Context, cc *genai.ClientConfig) (*Client, error) {
471+
if err := validateClientConfig(cc); err != nil {
472+
return nil, err
473+
}
464474
ac, err := genai.NewInternalAPIClient(ctx, cc)
465475
if err != nil {
466476
return nil, err

vertexai/genai/client_test.go

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import (
2828
"testing"
2929

3030
"google.golang.org/api/iterator"
31+
"google.golang.org/genai"
3132
)
3233

3334
const defaultModel = "gemini-1.0-pro"
@@ -739,3 +740,40 @@ func printResponse(resp *GenerateContentResponse) {
739740
}
740741
fmt.Println("---")
741742
}
743+
744+
func TestNewGenAIClient(t *testing.T) {
745+
ctx := context.Background()
746+
for _, test := range []struct {
747+
name string
748+
cc *genai.ClientConfig
749+
}{
750+
{name: "nil config", cc: nil},
751+
{name: "empty config", cc: &genai.ClientConfig{}},
752+
} {
753+
t.Run(test.name, func(t *testing.T) {
754+
client, err := NewGenAIClient(ctx, test.cc)
755+
if err != nil {
756+
t.Fatalf("NewGenAIClient() failed unexpectedly, err: %v", err)
757+
}
758+
if client == nil {
759+
t.Error("client must not be nil")
760+
}
761+
})
762+
}
763+
}
764+
765+
func TestNewGenAIClientErrors(t *testing.T) {
766+
ctx := context.Background()
767+
for _, test := range []struct {
768+
name string
769+
cc *genai.ClientConfig
770+
}{
771+
{name: "gemini backend", cc: &genai.ClientConfig{Backend: genai.BackendGeminiAPI}},
772+
} {
773+
t.Run(test.name, func(t *testing.T) {
774+
if _, err := NewGenAIClient(ctx, test.cc); err == nil {
775+
t.Error("wants error, but got nil")
776+
}
777+
})
778+
}
779+
}

0 commit comments

Comments
 (0)