mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-02-15 01:25:13 -05:00
## Description --- Update out stream in `stdio` transport protocol. **Before:** When `stdio=true`, log levels will automatically get set to WARN (even if user explicitly set it as DEBUG or INFO). This is because the DEBUG and INFO logger send logs to outStream. In `stdio`, server MUST NOT write anything to `stdout` that is not a valid MCP message. MCP specifications mentioned that "The server MAY write UTF-8 strings to its standard error (stderr) for logging purposes. Client MAY capture, forward, or ignore this logging." **After:** When `stdio=true`, logger's DEBUG and INFO logger send logs to outStream. This will still respect the `log-levels` set by user (default as INFO).
1600 lines
46 KiB
Go
1600 lines
46 KiB
Go
// Copyright 2024 Google LLC
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package cmd
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
_ "embed"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"path"
|
|
"path/filepath"
|
|
"regexp"
|
|
"runtime"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
|
|
"github.com/googleapis/genai-toolbox/internal/auth/google"
|
|
"github.com/googleapis/genai-toolbox/internal/log"
|
|
"github.com/googleapis/genai-toolbox/internal/prebuiltconfigs"
|
|
"github.com/googleapis/genai-toolbox/internal/server"
|
|
cloudsqlpgsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
|
|
httpsrc "github.com/googleapis/genai-toolbox/internal/sources/http"
|
|
"github.com/googleapis/genai-toolbox/internal/telemetry"
|
|
"github.com/googleapis/genai-toolbox/internal/testutils"
|
|
"github.com/googleapis/genai-toolbox/internal/tools"
|
|
"github.com/googleapis/genai-toolbox/internal/tools/http"
|
|
"github.com/googleapis/genai-toolbox/internal/tools/postgres/postgressql"
|
|
"github.com/googleapis/genai-toolbox/internal/util"
|
|
"github.com/spf13/cobra"
|
|
)
|
|
|
|
func withDefaults(c server.ServerConfig) server.ServerConfig {
|
|
data, _ := os.ReadFile("version.txt")
|
|
version := strings.TrimSpace(string(data)) // Preserving 'data', new var for clarity
|
|
c.Version = version + "+" + strings.Join([]string{"dev", runtime.GOOS, runtime.GOARCH}, ".")
|
|
|
|
if c.Address == "" {
|
|
c.Address = "127.0.0.1"
|
|
}
|
|
if c.Port == 0 {
|
|
c.Port = 5000
|
|
}
|
|
if c.TelemetryServiceName == "" {
|
|
c.TelemetryServiceName = "toolbox"
|
|
}
|
|
return c
|
|
}
|
|
|
|
func invokeCommand(args []string) (*Command, string, error) {
|
|
c := NewCommand()
|
|
|
|
// Keep the test output quiet
|
|
c.SilenceUsage = true
|
|
c.SilenceErrors = true
|
|
|
|
// Capture output
|
|
buf := new(bytes.Buffer)
|
|
c.SetOut(buf)
|
|
c.SetErr(buf)
|
|
c.SetArgs(args)
|
|
|
|
// Disable execute behavior
|
|
c.RunE = func(*cobra.Command, []string) error {
|
|
return nil
|
|
}
|
|
|
|
err := c.Execute()
|
|
|
|
return c, buf.String(), err
|
|
}
|
|
|
|
func TestVersion(t *testing.T) {
|
|
data, err := os.ReadFile("version.txt")
|
|
if err != nil {
|
|
t.Fatalf("failed to read version.txt: %v", err)
|
|
}
|
|
want := strings.TrimSpace(string(data))
|
|
|
|
_, got, err := invokeCommand([]string{"--version"})
|
|
if err != nil {
|
|
t.Fatalf("error invoking command: %s", err)
|
|
}
|
|
|
|
if !strings.Contains(got, want) {
|
|
t.Errorf("cli did not return correct version: want %q, got %q", want, got)
|
|
}
|
|
}
|
|
|
|
func TestServerConfigFlags(t *testing.T) {
|
|
tcs := []struct {
|
|
desc string
|
|
args []string
|
|
want server.ServerConfig
|
|
}{
|
|
{
|
|
desc: "default values",
|
|
args: []string{},
|
|
want: withDefaults(server.ServerConfig{}),
|
|
},
|
|
{
|
|
desc: "address short",
|
|
args: []string{"-a", "127.0.1.1"},
|
|
want: withDefaults(server.ServerConfig{
|
|
Address: "127.0.1.1",
|
|
}),
|
|
},
|
|
{
|
|
desc: "address long",
|
|
args: []string{"--address", "0.0.0.0"},
|
|
want: withDefaults(server.ServerConfig{
|
|
Address: "0.0.0.0",
|
|
}),
|
|
},
|
|
{
|
|
desc: "port short",
|
|
args: []string{"-p", "5052"},
|
|
want: withDefaults(server.ServerConfig{
|
|
Port: 5052,
|
|
}),
|
|
},
|
|
{
|
|
desc: "port long",
|
|
args: []string{"--port", "5050"},
|
|
want: withDefaults(server.ServerConfig{
|
|
Port: 5050,
|
|
}),
|
|
},
|
|
{
|
|
desc: "logging format",
|
|
args: []string{"--logging-format", "JSON"},
|
|
want: withDefaults(server.ServerConfig{
|
|
LoggingFormat: "JSON",
|
|
}),
|
|
},
|
|
{
|
|
desc: "debug logs",
|
|
args: []string{"--log-level", "WARN"},
|
|
want: withDefaults(server.ServerConfig{
|
|
LogLevel: "WARN",
|
|
}),
|
|
},
|
|
{
|
|
desc: "telemetry gcp",
|
|
args: []string{"--telemetry-gcp"},
|
|
want: withDefaults(server.ServerConfig{
|
|
TelemetryGCP: true,
|
|
}),
|
|
},
|
|
{
|
|
desc: "telemetry otlp",
|
|
args: []string{"--telemetry-otlp", "http://127.0.0.1:4553"},
|
|
want: withDefaults(server.ServerConfig{
|
|
TelemetryOTLP: "http://127.0.0.1:4553",
|
|
}),
|
|
},
|
|
{
|
|
desc: "telemetry service name",
|
|
args: []string{"--telemetry-service-name", "toolbox-custom"},
|
|
want: withDefaults(server.ServerConfig{
|
|
TelemetryServiceName: "toolbox-custom",
|
|
}),
|
|
},
|
|
{
|
|
desc: "stdio",
|
|
args: []string{"--stdio"},
|
|
want: withDefaults(server.ServerConfig{
|
|
Stdio: true,
|
|
}),
|
|
},
|
|
{
|
|
desc: "disable reload",
|
|
args: []string{"--disable-reload"},
|
|
want: withDefaults(server.ServerConfig{
|
|
DisableReload: true,
|
|
}),
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.desc, func(t *testing.T) {
|
|
c, _, err := invokeCommand(tc.args)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error invoking command: %s", err)
|
|
}
|
|
|
|
if !cmp.Equal(c.cfg, tc.want) {
|
|
t.Fatalf("got %v, want %v", c.cfg, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestParseEnv(t *testing.T) {
|
|
tcs := []struct {
|
|
desc string
|
|
env map[string]string
|
|
in string
|
|
want string
|
|
err bool
|
|
errString string
|
|
}{
|
|
{
|
|
desc: "without default without env",
|
|
in: "${FOO}",
|
|
want: "",
|
|
err: true,
|
|
errString: `environment variable not found: "FOO"`,
|
|
},
|
|
{
|
|
desc: "without default with env",
|
|
env: map[string]string{
|
|
"FOO": "bar",
|
|
},
|
|
in: "${FOO}",
|
|
want: "bar",
|
|
},
|
|
{
|
|
desc: "with empty default",
|
|
in: "${FOO:}",
|
|
want: "",
|
|
},
|
|
{
|
|
desc: "with default",
|
|
in: "${FOO:bar}",
|
|
want: "bar",
|
|
},
|
|
{
|
|
desc: "with default with env",
|
|
env: map[string]string{
|
|
"FOO": "hello",
|
|
},
|
|
in: "${FOO:bar}",
|
|
want: "hello",
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.desc, func(t *testing.T) {
|
|
if tc.env != nil {
|
|
for k, v := range tc.env {
|
|
t.Setenv(k, v)
|
|
}
|
|
}
|
|
got, err := parseEnv(tc.in)
|
|
if tc.err {
|
|
if err == nil {
|
|
t.Fatalf("expected error not found")
|
|
}
|
|
if tc.errString != err.Error() {
|
|
t.Fatalf("incorrect error string: got %s, want %s", err, tc.errString)
|
|
}
|
|
}
|
|
if tc.want != got {
|
|
t.Fatalf("unexpected want: got %s, want %s", got, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestToolFileFlag(t *testing.T) {
|
|
tcs := []struct {
|
|
desc string
|
|
args []string
|
|
want string
|
|
}{
|
|
{
|
|
desc: "default value",
|
|
args: []string{},
|
|
want: "",
|
|
},
|
|
{
|
|
desc: "foo file",
|
|
args: []string{"--tools-file", "foo.yaml"},
|
|
want: "foo.yaml",
|
|
},
|
|
{
|
|
desc: "address long",
|
|
args: []string{"--tools-file", "bar.yaml"},
|
|
want: "bar.yaml",
|
|
},
|
|
{
|
|
desc: "deprecated flag",
|
|
args: []string{"--tools_file", "foo.yaml"},
|
|
want: "foo.yaml",
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.desc, func(t *testing.T) {
|
|
c, _, err := invokeCommand(tc.args)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error invoking command: %s", err)
|
|
}
|
|
if c.tools_file != tc.want {
|
|
t.Fatalf("got %v, want %v", c.cfg, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestToolsFilesFlag(t *testing.T) {
|
|
tcs := []struct {
|
|
desc string
|
|
args []string
|
|
want []string
|
|
}{
|
|
{
|
|
desc: "no value",
|
|
args: []string{},
|
|
want: []string{},
|
|
},
|
|
{
|
|
desc: "single file",
|
|
args: []string{"--tools-files", "foo.yaml"},
|
|
want: []string{"foo.yaml"},
|
|
},
|
|
{
|
|
desc: "multiple files",
|
|
args: []string{"--tools-files", "foo.yaml,bar.yaml"},
|
|
want: []string{"foo.yaml", "bar.yaml"},
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.desc, func(t *testing.T) {
|
|
c, _, err := invokeCommand(tc.args)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error invoking command: %s", err)
|
|
}
|
|
if diff := cmp.Diff(c.tools_files, tc.want); diff != "" {
|
|
t.Fatalf("got %v, want %v", c.tools_files, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestToolsFolderFlag(t *testing.T) {
|
|
tcs := []struct {
|
|
desc string
|
|
args []string
|
|
want string
|
|
}{
|
|
{
|
|
desc: "no value",
|
|
args: []string{},
|
|
want: "",
|
|
},
|
|
{
|
|
desc: "folder set",
|
|
args: []string{"--tools-folder", "test-folder"},
|
|
want: "test-folder",
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.desc, func(t *testing.T) {
|
|
c, _, err := invokeCommand(tc.args)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error invoking command: %s", err)
|
|
}
|
|
if c.tools_folder != tc.want {
|
|
t.Fatalf("got %v, want %v", c.tools_folder, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestPrebuiltFlag(t *testing.T) {
|
|
tcs := []struct {
|
|
desc string
|
|
args []string
|
|
want string
|
|
}{
|
|
{
|
|
desc: "default value",
|
|
args: []string{},
|
|
want: "",
|
|
},
|
|
{
|
|
desc: "custom pre built flag",
|
|
args: []string{"--tools-file", "alloydb"},
|
|
want: "alloydb",
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.desc, func(t *testing.T) {
|
|
c, _, err := invokeCommand(tc.args)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error invoking command: %s", err)
|
|
}
|
|
if c.tools_file != tc.want {
|
|
t.Fatalf("got %v, want %v", c.cfg, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFailServerConfigFlags(t *testing.T) {
|
|
tcs := []struct {
|
|
desc string
|
|
args []string
|
|
}{
|
|
{
|
|
desc: "logging format",
|
|
args: []string{"--logging-format", "fail"},
|
|
},
|
|
{
|
|
desc: "debug logs",
|
|
args: []string{"--log-level", "fail"},
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.desc, func(t *testing.T) {
|
|
_, _, err := invokeCommand(tc.args)
|
|
if err == nil {
|
|
t.Fatalf("expected an error, but got nil")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestDefaultLoggingFormat(t *testing.T) {
|
|
c, _, err := invokeCommand([]string{})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error invoking command: %s", err)
|
|
}
|
|
got := c.cfg.LoggingFormat.String()
|
|
want := "standard"
|
|
if got != want {
|
|
t.Fatalf("unexpected default logging format flag: got %v, want %v", got, want)
|
|
}
|
|
}
|
|
|
|
func TestDefaultLogLevel(t *testing.T) {
|
|
c, _, err := invokeCommand([]string{})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error invoking command: %s", err)
|
|
}
|
|
got := c.cfg.LogLevel.String()
|
|
want := "info"
|
|
if got != want {
|
|
t.Fatalf("unexpected default log level flag: got %v, want %v", got, want)
|
|
}
|
|
}
|
|
|
|
func TestParseToolFile(t *testing.T) {
|
|
ctx, err := testutils.ContextWithNewLogger()
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %s", err)
|
|
}
|
|
tcs := []struct {
|
|
description string
|
|
in string
|
|
wantToolsFile ToolsFile
|
|
}{
|
|
{
|
|
description: "basic example",
|
|
in: `
|
|
sources:
|
|
my-pg-instance:
|
|
kind: cloud-sql-postgres
|
|
project: my-project
|
|
region: my-region
|
|
instance: my-instance
|
|
database: my_db
|
|
user: my_user
|
|
password: my_pass
|
|
tools:
|
|
example_tool:
|
|
kind: postgres-sql
|
|
source: my-pg-instance
|
|
description: some description
|
|
statement: |
|
|
SELECT * FROM SQL_STATEMENT;
|
|
parameters:
|
|
- name: country
|
|
type: string
|
|
description: some description
|
|
toolsets:
|
|
example_toolset:
|
|
- example_tool
|
|
`,
|
|
wantToolsFile: ToolsFile{
|
|
Sources: server.SourceConfigs{
|
|
"my-pg-instance": cloudsqlpgsrc.Config{
|
|
Name: "my-pg-instance",
|
|
Kind: cloudsqlpgsrc.SourceKind,
|
|
Project: "my-project",
|
|
Region: "my-region",
|
|
Instance: "my-instance",
|
|
IPType: "public",
|
|
Database: "my_db",
|
|
User: "my_user",
|
|
Password: "my_pass",
|
|
},
|
|
},
|
|
Tools: server.ToolConfigs{
|
|
"example_tool": postgressql.Config{
|
|
Name: "example_tool",
|
|
Kind: "postgres-sql",
|
|
Source: "my-pg-instance",
|
|
Description: "some description",
|
|
Statement: "SELECT * FROM SQL_STATEMENT;\n",
|
|
Parameters: []tools.Parameter{
|
|
tools.NewStringParameter("country", "some description"),
|
|
},
|
|
AuthRequired: []string{},
|
|
},
|
|
},
|
|
Toolsets: server.ToolsetConfigs{
|
|
"example_toolset": tools.ToolsetConfig{
|
|
Name: "example_toolset",
|
|
ToolNames: []string{"example_tool"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.description, func(t *testing.T) {
|
|
toolsFile, err := parseToolsFile(ctx, testutils.FormatYaml(tc.in))
|
|
if err != nil {
|
|
t.Fatalf("failed to parse input: %v", err)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.Sources, toolsFile.Sources); diff != "" {
|
|
t.Fatalf("incorrect sources parse: diff %v", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.AuthServices, toolsFile.AuthServices); diff != "" {
|
|
t.Fatalf("incorrect authServices parse: diff %v", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.Tools, toolsFile.Tools); diff != "" {
|
|
t.Fatalf("incorrect tools parse: diff %v", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.Toolsets, toolsFile.Toolsets); diff != "" {
|
|
t.Fatalf("incorrect tools parse: diff %v", diff)
|
|
}
|
|
})
|
|
}
|
|
|
|
}
|
|
|
|
func TestParseToolFileWithAuth(t *testing.T) {
|
|
ctx, err := testutils.ContextWithNewLogger()
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %s", err)
|
|
}
|
|
tcs := []struct {
|
|
description string
|
|
in string
|
|
wantToolsFile ToolsFile
|
|
}{
|
|
{
|
|
description: "basic example",
|
|
in: `
|
|
sources:
|
|
my-pg-instance:
|
|
kind: cloud-sql-postgres
|
|
project: my-project
|
|
region: my-region
|
|
instance: my-instance
|
|
database: my_db
|
|
user: my_user
|
|
password: my_pass
|
|
authServices:
|
|
my-google-service:
|
|
kind: google
|
|
clientId: my-client-id
|
|
other-google-service:
|
|
kind: google
|
|
clientId: other-client-id
|
|
|
|
tools:
|
|
example_tool:
|
|
kind: postgres-sql
|
|
source: my-pg-instance
|
|
description: some description
|
|
statement: |
|
|
SELECT * FROM SQL_STATEMENT;
|
|
parameters:
|
|
- name: country
|
|
type: string
|
|
description: some description
|
|
- name: id
|
|
type: integer
|
|
description: user id
|
|
authServices:
|
|
- name: my-google-service
|
|
field: user_id
|
|
- name: email
|
|
type: string
|
|
description: user email
|
|
authServices:
|
|
- name: my-google-service
|
|
field: email
|
|
- name: other-google-service
|
|
field: other_email
|
|
|
|
toolsets:
|
|
example_toolset:
|
|
- example_tool
|
|
`,
|
|
wantToolsFile: ToolsFile{
|
|
Sources: server.SourceConfigs{
|
|
"my-pg-instance": cloudsqlpgsrc.Config{
|
|
Name: "my-pg-instance",
|
|
Kind: cloudsqlpgsrc.SourceKind,
|
|
Project: "my-project",
|
|
Region: "my-region",
|
|
Instance: "my-instance",
|
|
IPType: "public",
|
|
Database: "my_db",
|
|
User: "my_user",
|
|
Password: "my_pass",
|
|
},
|
|
},
|
|
AuthServices: server.AuthServiceConfigs{
|
|
"my-google-service": google.Config{
|
|
Name: "my-google-service",
|
|
Kind: google.AuthServiceKind,
|
|
ClientID: "my-client-id",
|
|
},
|
|
"other-google-service": google.Config{
|
|
Name: "other-google-service",
|
|
Kind: google.AuthServiceKind,
|
|
ClientID: "other-client-id",
|
|
},
|
|
},
|
|
Tools: server.ToolConfigs{
|
|
"example_tool": postgressql.Config{
|
|
Name: "example_tool",
|
|
Kind: "postgres-sql",
|
|
Source: "my-pg-instance",
|
|
Description: "some description",
|
|
Statement: "SELECT * FROM SQL_STATEMENT;\n",
|
|
AuthRequired: []string{},
|
|
Parameters: []tools.Parameter{
|
|
tools.NewStringParameter("country", "some description"),
|
|
tools.NewIntParameterWithAuth("id", "user id", []tools.ParamAuthService{{Name: "my-google-service", Field: "user_id"}}),
|
|
tools.NewStringParameterWithAuth("email", "user email", []tools.ParamAuthService{{Name: "my-google-service", Field: "email"}, {Name: "other-google-service", Field: "other_email"}}),
|
|
},
|
|
},
|
|
},
|
|
Toolsets: server.ToolsetConfigs{
|
|
"example_toolset": tools.ToolsetConfig{
|
|
Name: "example_toolset",
|
|
ToolNames: []string{"example_tool"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
description: "basic example with authSources",
|
|
in: `
|
|
sources:
|
|
my-pg-instance:
|
|
kind: cloud-sql-postgres
|
|
project: my-project
|
|
region: my-region
|
|
instance: my-instance
|
|
database: my_db
|
|
user: my_user
|
|
password: my_pass
|
|
authSources:
|
|
my-google-service:
|
|
kind: google
|
|
clientId: my-client-id
|
|
other-google-service:
|
|
kind: google
|
|
clientId: other-client-id
|
|
|
|
tools:
|
|
example_tool:
|
|
kind: postgres-sql
|
|
source: my-pg-instance
|
|
description: some description
|
|
statement: |
|
|
SELECT * FROM SQL_STATEMENT;
|
|
parameters:
|
|
- name: country
|
|
type: string
|
|
description: some description
|
|
- name: id
|
|
type: integer
|
|
description: user id
|
|
authSources:
|
|
- name: my-google-service
|
|
field: user_id
|
|
- name: email
|
|
type: string
|
|
description: user email
|
|
authSources:
|
|
- name: my-google-service
|
|
field: email
|
|
- name: other-google-service
|
|
field: other_email
|
|
|
|
toolsets:
|
|
example_toolset:
|
|
- example_tool
|
|
`,
|
|
wantToolsFile: ToolsFile{
|
|
Sources: server.SourceConfigs{
|
|
"my-pg-instance": cloudsqlpgsrc.Config{
|
|
Name: "my-pg-instance",
|
|
Kind: cloudsqlpgsrc.SourceKind,
|
|
Project: "my-project",
|
|
Region: "my-region",
|
|
Instance: "my-instance",
|
|
IPType: "public",
|
|
Database: "my_db",
|
|
User: "my_user",
|
|
Password: "my_pass",
|
|
},
|
|
},
|
|
AuthSources: server.AuthServiceConfigs{
|
|
"my-google-service": google.Config{
|
|
Name: "my-google-service",
|
|
Kind: google.AuthServiceKind,
|
|
ClientID: "my-client-id",
|
|
},
|
|
"other-google-service": google.Config{
|
|
Name: "other-google-service",
|
|
Kind: google.AuthServiceKind,
|
|
ClientID: "other-client-id",
|
|
},
|
|
},
|
|
Tools: server.ToolConfigs{
|
|
"example_tool": postgressql.Config{
|
|
Name: "example_tool",
|
|
Kind: "postgres-sql",
|
|
Source: "my-pg-instance",
|
|
Description: "some description",
|
|
Statement: "SELECT * FROM SQL_STATEMENT;\n",
|
|
AuthRequired: []string{},
|
|
Parameters: []tools.Parameter{
|
|
tools.NewStringParameter("country", "some description"),
|
|
tools.NewIntParameterWithAuth("id", "user id", []tools.ParamAuthService{{Name: "my-google-service", Field: "user_id"}}),
|
|
tools.NewStringParameterWithAuth("email", "user email", []tools.ParamAuthService{{Name: "my-google-service", Field: "email"}, {Name: "other-google-service", Field: "other_email"}}),
|
|
},
|
|
},
|
|
},
|
|
Toolsets: server.ToolsetConfigs{
|
|
"example_toolset": tools.ToolsetConfig{
|
|
Name: "example_toolset",
|
|
ToolNames: []string{"example_tool"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
description: "basic example with authRequired",
|
|
in: `
|
|
sources:
|
|
my-pg-instance:
|
|
kind: cloud-sql-postgres
|
|
project: my-project
|
|
region: my-region
|
|
instance: my-instance
|
|
database: my_db
|
|
user: my_user
|
|
password: my_pass
|
|
authServices:
|
|
my-google-service:
|
|
kind: google
|
|
clientId: my-client-id
|
|
other-google-service:
|
|
kind: google
|
|
clientId: other-client-id
|
|
|
|
tools:
|
|
example_tool:
|
|
kind: postgres-sql
|
|
source: my-pg-instance
|
|
description: some description
|
|
statement: |
|
|
SELECT * FROM SQL_STATEMENT;
|
|
authRequired:
|
|
- my-google-service
|
|
parameters:
|
|
- name: country
|
|
type: string
|
|
description: some description
|
|
- name: id
|
|
type: integer
|
|
description: user id
|
|
authServices:
|
|
- name: my-google-service
|
|
field: user_id
|
|
- name: email
|
|
type: string
|
|
description: user email
|
|
authServices:
|
|
- name: my-google-service
|
|
field: email
|
|
- name: other-google-service
|
|
field: other_email
|
|
|
|
toolsets:
|
|
example_toolset:
|
|
- example_tool
|
|
`,
|
|
wantToolsFile: ToolsFile{
|
|
Sources: server.SourceConfigs{
|
|
"my-pg-instance": cloudsqlpgsrc.Config{
|
|
Name: "my-pg-instance",
|
|
Kind: cloudsqlpgsrc.SourceKind,
|
|
Project: "my-project",
|
|
Region: "my-region",
|
|
Instance: "my-instance",
|
|
IPType: "public",
|
|
Database: "my_db",
|
|
User: "my_user",
|
|
Password: "my_pass",
|
|
},
|
|
},
|
|
AuthServices: server.AuthServiceConfigs{
|
|
"my-google-service": google.Config{
|
|
Name: "my-google-service",
|
|
Kind: google.AuthServiceKind,
|
|
ClientID: "my-client-id",
|
|
},
|
|
"other-google-service": google.Config{
|
|
Name: "other-google-service",
|
|
Kind: google.AuthServiceKind,
|
|
ClientID: "other-client-id",
|
|
},
|
|
},
|
|
Tools: server.ToolConfigs{
|
|
"example_tool": postgressql.Config{
|
|
Name: "example_tool",
|
|
Kind: "postgres-sql",
|
|
Source: "my-pg-instance",
|
|
Description: "some description",
|
|
Statement: "SELECT * FROM SQL_STATEMENT;\n",
|
|
AuthRequired: []string{"my-google-service"},
|
|
Parameters: []tools.Parameter{
|
|
tools.NewStringParameter("country", "some description"),
|
|
tools.NewIntParameterWithAuth("id", "user id", []tools.ParamAuthService{{Name: "my-google-service", Field: "user_id"}}),
|
|
tools.NewStringParameterWithAuth("email", "user email", []tools.ParamAuthService{{Name: "my-google-service", Field: "email"}, {Name: "other-google-service", Field: "other_email"}}),
|
|
},
|
|
},
|
|
},
|
|
Toolsets: server.ToolsetConfigs{
|
|
"example_toolset": tools.ToolsetConfig{
|
|
Name: "example_toolset",
|
|
ToolNames: []string{"example_tool"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.description, func(t *testing.T) {
|
|
toolsFile, err := parseToolsFile(ctx, testutils.FormatYaml(tc.in))
|
|
if err != nil {
|
|
t.Fatalf("failed to parse input: %v", err)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.Sources, toolsFile.Sources); diff != "" {
|
|
t.Fatalf("incorrect sources parse: diff %v", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.AuthServices, toolsFile.AuthServices); diff != "" {
|
|
t.Fatalf("incorrect authServices parse: diff %v", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.Tools, toolsFile.Tools); diff != "" {
|
|
t.Fatalf("incorrect tools parse: diff %v", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.Toolsets, toolsFile.Toolsets); diff != "" {
|
|
t.Fatalf("incorrect tools parse: diff %v", diff)
|
|
}
|
|
})
|
|
}
|
|
|
|
}
|
|
|
|
func TestEnvVarReplacement(t *testing.T) {
|
|
ctx, err := testutils.ContextWithNewLogger()
|
|
t.Setenv("TestHeader", "ACTUAL_HEADER")
|
|
t.Setenv("API_KEY", "ACTUAL_API_KEY")
|
|
t.Setenv("clientId", "ACTUAL_CLIENT_ID")
|
|
t.Setenv("clientId2", "ACTUAL_CLIENT_ID_2")
|
|
t.Setenv("toolset_name", "ACTUAL_TOOLSET_NAME")
|
|
t.Setenv("cat_string", "cat")
|
|
t.Setenv("food_string", "food")
|
|
t.Setenv("TestHeader", "ACTUAL_HEADER")
|
|
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %s", err)
|
|
}
|
|
tcs := []struct {
|
|
description string
|
|
in string
|
|
wantToolsFile ToolsFile
|
|
}{
|
|
{
|
|
description: "file with env var example",
|
|
in: `
|
|
sources:
|
|
my-http-instance:
|
|
kind: http
|
|
baseUrl: http://test_server/
|
|
timeout: 10s
|
|
headers:
|
|
Authorization: ${TestHeader}
|
|
queryParams:
|
|
api-key: ${API_KEY}
|
|
authServices:
|
|
my-google-service:
|
|
kind: google
|
|
clientId: ${clientId}
|
|
other-google-service:
|
|
kind: google
|
|
clientId: ${clientId2}
|
|
|
|
tools:
|
|
example_tool:
|
|
kind: http
|
|
source: my-instance
|
|
method: GET
|
|
path: "search?name=alice&pet=${cat_string}"
|
|
description: some description
|
|
authRequired:
|
|
- my-google-auth-service
|
|
- other-auth-service
|
|
queryParams:
|
|
- name: country
|
|
type: string
|
|
description: some description
|
|
authServices:
|
|
- name: my-google-auth-service
|
|
field: user_id
|
|
- name: other-auth-service
|
|
field: user_id
|
|
requestBody: |
|
|
{
|
|
"age": {{.age}},
|
|
"city": "{{.city}}",
|
|
"food": "${food_string}",
|
|
"other": "$OTHER"
|
|
}
|
|
bodyParams:
|
|
- name: age
|
|
type: integer
|
|
description: age num
|
|
- name: city
|
|
type: string
|
|
description: city string
|
|
headers:
|
|
Authorization: API_KEY
|
|
Content-Type: application/json
|
|
headerParams:
|
|
- name: Language
|
|
type: string
|
|
description: language string
|
|
|
|
toolsets:
|
|
${toolset_name}:
|
|
- example_tool
|
|
`,
|
|
wantToolsFile: ToolsFile{
|
|
Sources: server.SourceConfigs{
|
|
"my-http-instance": httpsrc.Config{
|
|
Name: "my-http-instance",
|
|
Kind: httpsrc.SourceKind,
|
|
BaseURL: "http://test_server/",
|
|
Timeout: "10s",
|
|
DefaultHeaders: map[string]string{"Authorization": "ACTUAL_HEADER"},
|
|
QueryParams: map[string]string{"api-key": "ACTUAL_API_KEY"},
|
|
},
|
|
},
|
|
AuthServices: server.AuthServiceConfigs{
|
|
"my-google-service": google.Config{
|
|
Name: "my-google-service",
|
|
Kind: google.AuthServiceKind,
|
|
ClientID: "ACTUAL_CLIENT_ID",
|
|
},
|
|
"other-google-service": google.Config{
|
|
Name: "other-google-service",
|
|
Kind: google.AuthServiceKind,
|
|
ClientID: "ACTUAL_CLIENT_ID_2",
|
|
},
|
|
},
|
|
Tools: server.ToolConfigs{
|
|
"example_tool": http.Config{
|
|
Name: "example_tool",
|
|
Kind: "http",
|
|
Source: "my-instance",
|
|
Method: "GET",
|
|
Path: "search?name=alice&pet=cat",
|
|
Description: "some description",
|
|
AuthRequired: []string{"my-google-auth-service", "other-auth-service"},
|
|
QueryParams: []tools.Parameter{
|
|
tools.NewStringParameterWithAuth("country", "some description",
|
|
[]tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_id"},
|
|
{Name: "other-auth-service", Field: "user_id"}}),
|
|
},
|
|
RequestBody: `{
|
|
"age": {{.age}},
|
|
"city": "{{.city}}",
|
|
"food": "food",
|
|
"other": "$OTHER"
|
|
}
|
|
`,
|
|
BodyParams: []tools.Parameter{tools.NewIntParameter("age", "age num"), tools.NewStringParameter("city", "city string")},
|
|
Headers: map[string]string{"Authorization": "API_KEY", "Content-Type": "application/json"},
|
|
HeaderParams: []tools.Parameter{tools.NewStringParameter("Language", "language string")},
|
|
},
|
|
},
|
|
Toolsets: server.ToolsetConfigs{
|
|
"ACTUAL_TOOLSET_NAME": tools.ToolsetConfig{
|
|
Name: "ACTUAL_TOOLSET_NAME",
|
|
ToolNames: []string{"example_tool"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.description, func(t *testing.T) {
|
|
toolsFile, err := parseToolsFile(ctx, testutils.FormatYaml(tc.in))
|
|
if err != nil {
|
|
t.Fatalf("failed to parse input: %v", err)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.Sources, toolsFile.Sources); diff != "" {
|
|
t.Fatalf("incorrect sources parse: diff %v", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.AuthServices, toolsFile.AuthServices); diff != "" {
|
|
t.Fatalf("incorrect authServices parse: diff %v", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.Tools, toolsFile.Tools); diff != "" {
|
|
t.Fatalf("incorrect tools parse: diff %v", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolsFile.Toolsets, toolsFile.Toolsets); diff != "" {
|
|
t.Fatalf("incorrect tools parse: diff %v", diff)
|
|
}
|
|
})
|
|
}
|
|
|
|
}
|
|
|
|
// normalizeFilepaths is a helper function to allow same filepath formats for Mac and Windows.
|
|
// this prevents needing multiple "want" cases for TestResolveWatcherInputs
|
|
func normalizeFilepaths(m map[string]bool) map[string]bool {
|
|
newMap := make(map[string]bool)
|
|
for k, v := range m {
|
|
newMap[filepath.ToSlash(k)] = v
|
|
}
|
|
return newMap
|
|
}
|
|
|
|
func TestResolveWatcherInputs(t *testing.T) {
|
|
tcs := []struct {
|
|
description string
|
|
toolsFile string
|
|
toolsFiles []string
|
|
toolsFolder string
|
|
wantWatchDirs map[string]bool
|
|
wantWatchedFiles map[string]bool
|
|
}{
|
|
{
|
|
description: "single tools file",
|
|
toolsFile: "tools_folder/example_tools.yaml",
|
|
toolsFiles: []string{},
|
|
toolsFolder: "",
|
|
wantWatchDirs: map[string]bool{"tools_folder": true},
|
|
wantWatchedFiles: map[string]bool{"tools_folder/example_tools.yaml": true},
|
|
},
|
|
{
|
|
description: "default tools file (root dir)",
|
|
toolsFile: "tools.yaml",
|
|
toolsFiles: []string{},
|
|
toolsFolder: "",
|
|
wantWatchDirs: map[string]bool{".": true},
|
|
wantWatchedFiles: map[string]bool{"tools.yaml": true},
|
|
},
|
|
{
|
|
description: "multiple files in different folders",
|
|
toolsFile: "",
|
|
toolsFiles: []string{"tools_folder/example_tools.yaml", "tools_folder2/example_tools.yaml"},
|
|
toolsFolder: "",
|
|
wantWatchDirs: map[string]bool{"tools_folder": true, "tools_folder2": true},
|
|
wantWatchedFiles: map[string]bool{
|
|
"tools_folder/example_tools.yaml": true,
|
|
"tools_folder2/example_tools.yaml": true,
|
|
},
|
|
},
|
|
{
|
|
description: "multiple files in same folder",
|
|
toolsFile: "",
|
|
toolsFiles: []string{"tools_folder/example_tools.yaml", "tools_folder/example_tools2.yaml"},
|
|
toolsFolder: "",
|
|
wantWatchDirs: map[string]bool{"tools_folder": true},
|
|
wantWatchedFiles: map[string]bool{
|
|
"tools_folder/example_tools.yaml": true,
|
|
"tools_folder/example_tools2.yaml": true,
|
|
},
|
|
},
|
|
{
|
|
description: "multiple files in different levels",
|
|
toolsFile: "",
|
|
toolsFiles: []string{
|
|
"tools_folder/example_tools.yaml",
|
|
"tools_folder/special_tools/example_tools2.yaml"},
|
|
toolsFolder: "",
|
|
wantWatchDirs: map[string]bool{"tools_folder": true, "tools_folder/special_tools": true},
|
|
wantWatchedFiles: map[string]bool{
|
|
"tools_folder/example_tools.yaml": true,
|
|
"tools_folder/special_tools/example_tools2.yaml": true,
|
|
},
|
|
},
|
|
{
|
|
description: "tools folder",
|
|
toolsFile: "",
|
|
toolsFiles: []string{},
|
|
toolsFolder: "tools_folder",
|
|
wantWatchDirs: map[string]bool{"tools_folder": true},
|
|
wantWatchedFiles: map[string]bool{},
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.description, func(t *testing.T) {
|
|
gotWatchDirs, gotWatchedFiles := resolveWatcherInputs(tc.toolsFile, tc.toolsFiles, tc.toolsFolder)
|
|
|
|
normalizedGotWatchDirs := normalizeFilepaths(gotWatchDirs)
|
|
normalizedGotWatchedFiles := normalizeFilepaths(gotWatchedFiles)
|
|
|
|
if diff := cmp.Diff(tc.wantWatchDirs, normalizedGotWatchDirs); diff != "" {
|
|
t.Errorf("incorrect watchDirs: diff %v", diff)
|
|
}
|
|
if diff := cmp.Diff(tc.wantWatchedFiles, normalizedGotWatchedFiles); diff != "" {
|
|
t.Errorf("incorrect watchedFiles: diff %v", diff)
|
|
}
|
|
|
|
})
|
|
}
|
|
}
|
|
|
|
// helper function for testing file detection in dynamic reloading
|
|
func tmpFileWithCleanup(content []byte) (string, func(), error) {
|
|
f, err := os.CreateTemp("", "*")
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
cleanup := func() { os.Remove(f.Name()) }
|
|
|
|
if _, err := f.Write(content); err != nil {
|
|
cleanup()
|
|
return "", nil, err
|
|
}
|
|
if err := f.Close(); err != nil {
|
|
cleanup()
|
|
return "", nil, err
|
|
}
|
|
return f.Name(), cleanup, err
|
|
}
|
|
|
|
func TestSingleEdit(t *testing.T) {
|
|
ctx, cancelCtx := context.WithTimeout(context.Background(), time.Minute)
|
|
defer cancelCtx()
|
|
|
|
pr, pw := io.Pipe()
|
|
defer pw.Close()
|
|
defer pr.Close()
|
|
|
|
fileToWatch, cleanup, err := tmpFileWithCleanup([]byte("initial content"))
|
|
if err != nil {
|
|
t.Fatalf("error editing tools file %s", err)
|
|
}
|
|
defer cleanup()
|
|
|
|
logger, err := log.NewStdLogger(pw, pw, "DEBUG")
|
|
if err != nil {
|
|
t.Fatalf("failed to setup logger %s", err)
|
|
}
|
|
ctx = util.WithLogger(ctx, logger)
|
|
|
|
instrumentation, err := telemetry.CreateTelemetryInstrumentation(versionString)
|
|
if err != nil {
|
|
t.Fatalf("failed to setup instrumentation %s", err)
|
|
}
|
|
ctx = util.WithInstrumentation(ctx, instrumentation)
|
|
|
|
mockServer := &server.Server{}
|
|
|
|
cleanFileToWatch := filepath.Clean(fileToWatch)
|
|
watchDir := filepath.Dir(cleanFileToWatch)
|
|
|
|
watchedFiles := map[string]bool{cleanFileToWatch: true}
|
|
watchDirs := map[string]bool{watchDir: true}
|
|
|
|
go watchChanges(ctx, watchDirs, watchedFiles, mockServer)
|
|
|
|
// escape backslash so regex doesn't fail on windows filepaths
|
|
regexEscapedPathFile := strings.ReplaceAll(cleanFileToWatch, `\`, `\\\\*\\`)
|
|
regexEscapedPathFile = path.Clean(regexEscapedPathFile)
|
|
|
|
regexEscapedPathDir := strings.ReplaceAll(watchDir, `\`, `\\\\*\\`)
|
|
regexEscapedPathDir = path.Clean(regexEscapedPathDir)
|
|
|
|
begunWatchingDir := regexp.MustCompile(fmt.Sprintf(`DEBUG "Added directory %s to watcher."`, regexEscapedPathDir))
|
|
_, err = testutils.WaitForString(ctx, begunWatchingDir, pr)
|
|
if err != nil {
|
|
t.Fatalf("timeout or error waiting for watcher to start: %s", err)
|
|
}
|
|
|
|
err = os.WriteFile(fileToWatch, []byte("modification"), 0777)
|
|
if err != nil {
|
|
t.Fatalf("error writing to file: %v", err)
|
|
}
|
|
|
|
// only check substring of DEBUG message due to some OS/editors firing different operations
|
|
detectedFileChange := regexp.MustCompile(fmt.Sprintf(`event detected in %s"`, regexEscapedPathFile))
|
|
_, err = testutils.WaitForString(ctx, detectedFileChange, pr)
|
|
if err != nil {
|
|
t.Fatalf("timeout or error waiting for file to detect write: %s", err)
|
|
}
|
|
}
|
|
|
|
func TestPrebuiltTools(t *testing.T) {
|
|
// Get prebuilt configs
|
|
alloydb_admin_config, _ := prebuiltconfigs.Get("alloydb-postgres-admin")
|
|
alloydb_config, _ := prebuiltconfigs.Get("alloydb-postgres")
|
|
bigquery_config, _ := prebuiltconfigs.Get("bigquery")
|
|
clickhouse_config, _ := prebuiltconfigs.Get("clickhouse")
|
|
cloudsqlpg_config, _ := prebuiltconfigs.Get("cloud-sql-postgres")
|
|
cloudsqlpg_admin_config, _ := prebuiltconfigs.Get("cloud-sql-postgres-admin")
|
|
cloudsqlmysql_config, _ := prebuiltconfigs.Get("cloud-sql-mysql")
|
|
cloudsqlmysql_admin_config, _ := prebuiltconfigs.Get("cloud-sql-mysql-admin")
|
|
cloudsqlmssql_config, _ := prebuiltconfigs.Get("cloud-sql-mssql")
|
|
cloudsqlmssql_admin_config, _ := prebuiltconfigs.Get("cloud-sql-mssql-admin")
|
|
dataplex_config, _ := prebuiltconfigs.Get("dataplex")
|
|
firestoreconfig, _ := prebuiltconfigs.Get("firestore")
|
|
mysql_config, _ := prebuiltconfigs.Get("mysql")
|
|
mssql_config, _ := prebuiltconfigs.Get("mssql")
|
|
looker_config, _ := prebuiltconfigs.Get("looker")
|
|
postgresconfig, _ := prebuiltconfigs.Get("postgres")
|
|
spanner_config, _ := prebuiltconfigs.Get("spanner")
|
|
spannerpg_config, _ := prebuiltconfigs.Get("spanner-postgres")
|
|
sqlite_config, _ := prebuiltconfigs.Get("sqlite")
|
|
neo4jconfig, _ := prebuiltconfigs.Get("neo4j")
|
|
alloydbobsvconfig, _ := prebuiltconfigs.Get("alloydb-postgres-observability")
|
|
cloudsqlpgobsvconfig, _ := prebuiltconfigs.Get("cloud-sql-postgres-observability")
|
|
cloudsqlmysqlobsvconfig, _ := prebuiltconfigs.Get("cloud-sql-mysql-observability")
|
|
cloudsqlmssqlobsvconfig, _ := prebuiltconfigs.Get("cloud-sql-mssql-observability")
|
|
|
|
// Set environment variables
|
|
t.Setenv("API_KEY", "your_api_key")
|
|
|
|
t.Setenv("BIGQUERY_PROJECT", "your_gcp_project_id")
|
|
t.Setenv("DATAPLEX_PROJECT", "your_gcp_project_id")
|
|
t.Setenv("FIRESTORE_PROJECT", "your_gcp_project_id")
|
|
t.Setenv("FIRESTORE_DATABASE", "your_firestore_db_name")
|
|
|
|
t.Setenv("SPANNER_PROJECT", "your_gcp_project_id")
|
|
t.Setenv("SPANNER_INSTANCE", "your_spanner_instance")
|
|
t.Setenv("SPANNER_DATABASE", "your_spanner_db")
|
|
|
|
t.Setenv("ALLOYDB_POSTGRES_PROJECT", "your_gcp_project_id")
|
|
t.Setenv("ALLOYDB_POSTGRES_REGION", "your_gcp_region")
|
|
t.Setenv("ALLOYDB_POSTGRES_CLUSTER", "your_alloydb_cluster")
|
|
t.Setenv("ALLOYDB_POSTGRES_INSTANCE", "your_alloydb_instance")
|
|
t.Setenv("ALLOYDB_POSTGRES_DATABASE", "your_alloydb_db")
|
|
t.Setenv("ALLOYDB_POSTGRES_USER", "your_alloydb_user")
|
|
t.Setenv("ALLOYDB_POSTGRES_PASSWORD", "your_alloydb_password")
|
|
|
|
t.Setenv("CLICKHOUSE_PROTOCOL", "your_clickhouse_protocol")
|
|
t.Setenv("CLICKHOUSE_DATABASE", "your_clickhouse_database")
|
|
t.Setenv("CLICKHOUSE_PASSWORD", "your_clickhouse_password")
|
|
t.Setenv("CLICKHOUSE_USER", "your_clickhouse_user")
|
|
t.Setenv("CLICKHOUSE_HOST", "your_clickhosue_host")
|
|
t.Setenv("CLICKHOUSE_PORT", "8123")
|
|
|
|
t.Setenv("CLOUD_SQL_POSTGRES_PROJECT", "your_pg_project")
|
|
t.Setenv("CLOUD_SQL_POSTGRES_INSTANCE", "your_pg_instance")
|
|
t.Setenv("CLOUD_SQL_POSTGRES_DATABASE", "your_pg_db")
|
|
t.Setenv("CLOUD_SQL_POSTGRES_REGION", "your_pg_region")
|
|
t.Setenv("CLOUD_SQL_POSTGRES_USER", "your_pg_user")
|
|
t.Setenv("CLOUD_SQL_POSTGRES_PASS", "your_pg_pass")
|
|
|
|
t.Setenv("CLOUD_SQL_MYSQL_PROJECT", "your_gcp_project_id")
|
|
t.Setenv("CLOUD_SQL_MYSQL_REGION", "your_gcp_region")
|
|
t.Setenv("CLOUD_SQL_MYSQL_INSTANCE", "your_instance")
|
|
t.Setenv("CLOUD_SQL_MYSQL_DATABASE", "your_cloudsql_mysql_db")
|
|
t.Setenv("CLOUD_SQL_MYSQL_USER", "your_cloudsql_mysql_user")
|
|
t.Setenv("CLOUD_SQL_MYSQL_PASSWORD", "your_cloudsql_mysql_password")
|
|
|
|
t.Setenv("CLOUD_SQL_MSSQL_PROJECT", "your_gcp_project_id")
|
|
t.Setenv("CLOUD_SQL_MSSQL_REGION", "your_gcp_region")
|
|
t.Setenv("CLOUD_SQL_MSSQL_INSTANCE", "your_cloudsql_mssql_instance")
|
|
t.Setenv("CLOUD_SQL_MSSQL_DATABASE", "your_cloudsql_mssql_db")
|
|
t.Setenv("CLOUD_SQL_MSSQL_IP_ADDRESS", "127.0.0.1")
|
|
t.Setenv("CLOUD_SQL_MSSQL_USER", "your_cloudsql_mssql_user")
|
|
t.Setenv("CLOUD_SQL_MSSQL_PASSWORD", "your_cloudsql_mssql_password")
|
|
t.Setenv("CLOUD_SQL_POSTGRES_PASSWORD", "your_cloudsql_pg_password")
|
|
|
|
t.Setenv("POSTGRES_HOST", "localhost")
|
|
t.Setenv("POSTGRES_PORT", "5432")
|
|
t.Setenv("POSTGRES_DATABASE", "your_postgres_db")
|
|
t.Setenv("POSTGRES_USER", "your_postgres_user")
|
|
t.Setenv("POSTGRES_PASSWORD", "your_postgres_password")
|
|
|
|
t.Setenv("MYSQL_HOST", "localhost")
|
|
t.Setenv("MYSQL_PORT", "3306")
|
|
t.Setenv("MYSQL_DATABASE", "your_mysql_db")
|
|
t.Setenv("MYSQL_USER", "your_mysql_user")
|
|
t.Setenv("MYSQL_PASSWORD", "your_mysql_password")
|
|
|
|
t.Setenv("MSSQL_HOST", "localhost")
|
|
t.Setenv("MSSQL_PORT", "1433")
|
|
t.Setenv("MSSQL_DATABASE", "your_mssql_db")
|
|
t.Setenv("MSSQL_USER", "your_mssql_user")
|
|
t.Setenv("MSSQL_PASSWORD", "your_mssql_password")
|
|
|
|
t.Setenv("LOOKER_BASE_URL", "https://your_company.looker.com")
|
|
t.Setenv("LOOKER_CLIENT_ID", "your_looker_client_id")
|
|
t.Setenv("LOOKER_CLIENT_SECRET", "your_looker_client_secret")
|
|
t.Setenv("LOOKER_VERIFY_SSL", "true")
|
|
|
|
t.Setenv("SQLITE_DATABASE", "test.db")
|
|
|
|
t.Setenv("NEO4J_URI", "bolt://localhost:7687")
|
|
t.Setenv("NEO4J_DATABASE", "neo4j")
|
|
t.Setenv("NEO4J_USERNAME", "your_neo4j_user")
|
|
t.Setenv("NEO4J_PASSWORD", "your_neo4j_password")
|
|
|
|
ctx, err := testutils.ContextWithNewLogger()
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %s", err)
|
|
}
|
|
tcs := []struct {
|
|
name string
|
|
in []byte
|
|
wantToolset server.ToolsetConfigs
|
|
}{
|
|
{
|
|
name: "alloydb postgres admin prebuilt tools",
|
|
in: alloydb_admin_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"alloydb_postgres_admin_tools": tools.ToolsetConfig{
|
|
Name: "alloydb_postgres_admin_tools",
|
|
ToolNames: []string{"create_cluster", "wait_for_operation", "create_instance", "list_clusters", "list_instances", "list_users", "create_user", "get_cluster", "get_instance", "get_user"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "cloudsql pg admin prebuilt tools",
|
|
in: cloudsqlpg_admin_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"cloud_sql_postgres_admin_tools": tools.ToolsetConfig{
|
|
Name: "cloud_sql_postgres_admin_tools",
|
|
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "cloudsql mysql admin prebuilt tools",
|
|
in: cloudsqlmysql_admin_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"cloud_sql_mysql_admin_tools": tools.ToolsetConfig{
|
|
Name: "cloud_sql_mysql_admin_tools",
|
|
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "cloudsql mssql admin prebuilt tools",
|
|
in: cloudsqlmssql_admin_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"cloud_sql_mssql_admin_tools": tools.ToolsetConfig{
|
|
Name: "cloud_sql_mssql_admin_tools",
|
|
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "alloydb prebuilt tools",
|
|
in: alloydb_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"alloydb_postgres_database_tools": tools.ToolsetConfig{
|
|
Name: "alloydb_postgres_database_tools",
|
|
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "bigquery prebuilt tools",
|
|
in: bigquery_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"bigquery_database_tools": tools.ToolsetConfig{
|
|
Name: "bigquery_database_tools",
|
|
ToolNames: []string{"analyze_contribution", "ask_data_insights", "execute_sql", "forecast", "get_dataset_info", "get_table_info", "list_dataset_ids", "list_table_ids", "search_catalog"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "clickhouse prebuilt tools",
|
|
in: clickhouse_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"clickhouse_database_tools": tools.ToolsetConfig{
|
|
Name: "clickhouse_database_tools",
|
|
ToolNames: []string{"execute_sql", "list_databases", "list_tables"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "cloudsqlpg prebuilt tools",
|
|
in: cloudsqlpg_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"cloud_sql_postgres_database_tools": tools.ToolsetConfig{
|
|
Name: "cloud_sql_postgres_database_tools",
|
|
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "cloudsqlmysql prebuilt tools",
|
|
in: cloudsqlmysql_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"cloud_sql_mysql_database_tools": tools.ToolsetConfig{
|
|
Name: "cloud_sql_mysql_database_tools",
|
|
ToolNames: []string{"execute_sql", "list_tables", "get_query_plan", "list_active_queries", "list_tables_missing_unique_indexes", "list_table_fragmentation"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "cloudsqlmssql prebuilt tools",
|
|
in: cloudsqlmssql_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"cloud_sql_mssql_database_tools": tools.ToolsetConfig{
|
|
Name: "cloud_sql_mssql_database_tools",
|
|
ToolNames: []string{"execute_sql", "list_tables"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "dataplex prebuilt tools",
|
|
in: dataplex_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"dataplex_tools": tools.ToolsetConfig{
|
|
Name: "dataplex_tools",
|
|
ToolNames: []string{"search_entries", "lookup_entry", "search_aspect_types"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "firestore prebuilt tools",
|
|
in: firestoreconfig,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"firestore_database_tools": tools.ToolsetConfig{
|
|
Name: "firestore_database_tools",
|
|
ToolNames: []string{"get_documents", "add_documents", "update_document", "list_collections", "delete_documents", "query_collection", "get_rules", "validate_rules"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "mysql prebuilt tools",
|
|
in: mysql_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"mysql_database_tools": tools.ToolsetConfig{
|
|
Name: "mysql_database_tools",
|
|
ToolNames: []string{"execute_sql", "list_tables", "get_query_plan", "list_active_queries", "list_tables_missing_unique_indexes", "list_table_fragmentation"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "mssql prebuilt tools",
|
|
in: mssql_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"mssql_database_tools": tools.ToolsetConfig{
|
|
Name: "mssql_database_tools",
|
|
ToolNames: []string{"execute_sql", "list_tables"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "looker prebuilt tools",
|
|
in: looker_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"looker_tools": tools.ToolsetConfig{
|
|
Name: "looker_tools",
|
|
ToolNames: []string{"get_models", "get_explores", "get_dimensions", "get_measures", "get_filters", "get_parameters", "query", "query_sql", "query_url", "get_looks", "run_look", "make_look", "get_dashboards", "make_dashboard", "add_dashboard_element"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "postgres prebuilt tools",
|
|
in: postgresconfig,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"postgres_database_tools": tools.ToolsetConfig{
|
|
Name: "postgres_database_tools",
|
|
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "spanner prebuilt tools",
|
|
in: spanner_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"spanner-database-tools": tools.ToolsetConfig{
|
|
Name: "spanner-database-tools",
|
|
ToolNames: []string{"execute_sql", "execute_sql_dql", "list_tables"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "spanner pg prebuilt tools",
|
|
in: spannerpg_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"spanner_postgres_database_tools": tools.ToolsetConfig{
|
|
Name: "spanner_postgres_database_tools",
|
|
ToolNames: []string{"execute_sql", "execute_sql_dql", "list_tables"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "sqlite prebuilt tools",
|
|
in: sqlite_config,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"sqlite_database_tools": tools.ToolsetConfig{
|
|
Name: "sqlite_database_tools",
|
|
ToolNames: []string{"execute_sql", "list_tables"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "neo4j prebuilt tools",
|
|
in: neo4jconfig,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"neo4j_database_tools": tools.ToolsetConfig{
|
|
Name: "neo4j_database_tools",
|
|
ToolNames: []string{"execute_cypher", "get_schema"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "alloydb postgres observability prebuilt tools",
|
|
in: alloydbobsvconfig,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"alloydb_postgres_cloud_monitoring_tools": tools.ToolsetConfig{
|
|
Name: "alloydb_postgres_cloud_monitoring_tools",
|
|
ToolNames: []string{"get_system_metrics", "get_query_metrics"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "cloudsql postgres observability prebuilt tools",
|
|
in: cloudsqlpgobsvconfig,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"cloud_sql_postgres_cloud_monitoring_tools": tools.ToolsetConfig{
|
|
Name: "cloud_sql_postgres_cloud_monitoring_tools",
|
|
ToolNames: []string{"get_system_metrics", "get_query_metrics"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "cloudsql mysql observability prebuilt tools",
|
|
in: cloudsqlmysqlobsvconfig,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"cloud_sql_mysql_cloud_monitoring_tools": tools.ToolsetConfig{
|
|
Name: "cloud_sql_mysql_cloud_monitoring_tools",
|
|
ToolNames: []string{"get_system_metrics", "get_query_metrics"},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
name: "cloudsql mssql observability prebuilt tools",
|
|
in: cloudsqlmssqlobsvconfig,
|
|
wantToolset: server.ToolsetConfigs{
|
|
"cloud_sql_mssql_cloud_monitoring_tools": tools.ToolsetConfig{
|
|
Name: "cloud_sql_mssql_cloud_monitoring_tools",
|
|
ToolNames: []string{"get_system_metrics"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
toolsFile, err := parseToolsFile(ctx, tc.in)
|
|
if err != nil {
|
|
t.Fatalf("failed to parse input: %v", err)
|
|
}
|
|
if diff := cmp.Diff(tc.wantToolset, toolsFile.Toolsets); diff != "" {
|
|
t.Fatalf("incorrect tools parse: diff %v", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|