// 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 tools_test import ( "bytes" "encoding/json" "math" "reflect" "testing" yaml "github.com/goccy/go-yaml" "github.com/google/go-cmp/cmp" "github.com/googleapis/genai-toolbox/internal/testutils" "github.com/googleapis/genai-toolbox/internal/tools" ) func TestParametersMarshal(t *testing.T) { ctx, err := testutils.ContextWithNewLogger() if err != nil { t.Fatalf("unexpected error: %s", err) } tcs := []struct { name string in []map[string]any want tools.Parameters }{ { name: "string", in: []map[string]any{ { "name": "my_string", "type": "string", "description": "this param is a string", }, }, want: tools.Parameters{ tools.NewStringParameter("my_string", "this param is a string"), }, }, { name: "int", in: []map[string]any{ { "name": "my_integer", "type": "integer", "description": "this param is an int", }, }, want: tools.Parameters{ tools.NewIntParameter("my_integer", "this param is an int"), }, }, { name: "float", in: []map[string]any{ { "name": "my_float", "type": "float", "description": "my param is a float", }, }, want: tools.Parameters{ tools.NewFloatParameter("my_float", "my param is a float"), }, }, { name: "bool", in: []map[string]any{ { "name": "my_bool", "type": "boolean", "description": "this param is a boolean", }, }, want: tools.Parameters{ tools.NewBooleanParameter("my_bool", "this param is a boolean"), }, }, { name: "string array", in: []map[string]any{ { "name": "my_array", "type": "array", "description": "this param is an array of strings", "items": map[string]string{ "name": "my_string", "type": "string", "description": "string item", }, }, }, want: tools.Parameters{ tools.NewArrayParameter("my_array", "this param is an array of strings", tools.NewStringParameter("my_string", "string item")), }, }, { name: "float array", in: []map[string]any{ { "name": "my_array", "type": "array", "description": "this param is an array of floats", "items": map[string]string{ "name": "my_float", "type": "float", "description": "float item", }, }, }, want: tools.Parameters{ tools.NewArrayParameter("my_array", "this param is an array of floats", tools.NewFloatParameter("my_float", "float item")), }, }, { name: "string default", in: []map[string]any{ { "name": "my_string", "type": "string", "default": "foo", "description": "this param is a string", }, }, want: tools.Parameters{ tools.NewStringParameterWithDefault("my_string", "foo", "this param is a string"), }, }, { name: "int default", in: []map[string]any{ { "name": "my_integer", "type": "integer", "default": 5, "description": "this param is an int", }, }, want: tools.Parameters{ tools.NewIntParameterWithDefault("my_integer", uint64(5), "this param is an int"), }, }, { name: "float default", in: []map[string]any{ { "name": "my_float", "type": "float", "default": 1.1, "description": "my param is a float", }, }, want: tools.Parameters{ tools.NewFloatParameterWithDefault("my_float", 1.1, "my param is a float"), }, }, { name: "bool default", in: []map[string]any{ { "name": "my_bool", "type": "boolean", "default": true, "description": "this param is a boolean", }, }, want: tools.Parameters{ tools.NewBooleanParameterWithDefault("my_bool", true, "this param is a boolean"), }, }, { name: "string array default", in: []map[string]any{ { "name": "my_array", "type": "array", "default": `["foo", "bar"]`, "description": "this param is an array of strings", "items": map[string]string{ "name": "my_string", "type": "string", "description": "string item", }, }, }, want: tools.Parameters{ tools.NewArrayParameterWithDefault("my_array", `["foo", "bar"]`, "this param is an array of strings", tools.NewStringParameter("my_string", "string item")), }, }, { name: "float array default", in: []map[string]any{ { "name": "my_array", "type": "array", "default": "[1.0, 1.1]", "description": "this param is an array of floats", "items": map[string]string{ "name": "my_float", "type": "float", "description": "float item", }, }, }, want: tools.Parameters{ tools.NewArrayParameterWithDefault("my_array", "[1.0, 1.1]", "this param is an array of floats", tools.NewFloatParameter("my_float", "float item")), }, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { var got tools.Parameters // parse map to bytes data, err := yaml.Marshal(tc.in) if err != nil { t.Fatalf("unable to marshal input to yaml: %s", err) } // parse bytes to object err = yaml.UnmarshalContext(ctx, data, &got) if err != nil { t.Fatalf("unable to unmarshal: %s", err) } if diff := cmp.Diff(tc.want, got); diff != "" { t.Fatalf("incorrect parse: diff %v", diff) } }) } } func TestAuthParametersMarshal(t *testing.T) { ctx, err := testutils.ContextWithNewLogger() if err != nil { t.Fatalf("unexpected error: %s", err) } authServices := []tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_id"}, {Name: "other-auth-service", Field: "user_id"}} tcs := []struct { name string in []map[string]any want tools.Parameters }{ { name: "string", in: []map[string]any{ { "name": "my_string", "type": "string", "description": "this param is a string", "authServices": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewStringParameterWithAuth("my_string", "this param is a string", authServices), }, }, { name: "string with authSources", in: []map[string]any{ { "name": "my_string", "type": "string", "description": "this param is a string", "authSources": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewStringParameterWithAuth("my_string", "this param is a string", authServices), }, }, { name: "int", in: []map[string]any{ { "name": "my_integer", "type": "integer", "description": "this param is an int", "authServices": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewIntParameterWithAuth("my_integer", "this param is an int", authServices), }, }, { name: "int with authSources", in: []map[string]any{ { "name": "my_integer", "type": "integer", "description": "this param is an int", "authSources": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewIntParameterWithAuth("my_integer", "this param is an int", authServices), }, }, { name: "float", in: []map[string]any{ { "name": "my_float", "type": "float", "description": "my param is a float", "authServices": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewFloatParameterWithAuth("my_float", "my param is a float", authServices), }, }, { name: "float with authSources", in: []map[string]any{ { "name": "my_float", "type": "float", "description": "my param is a float", "authSources": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewFloatParameterWithAuth("my_float", "my param is a float", authServices), }, }, { name: "bool", in: []map[string]any{ { "name": "my_bool", "type": "boolean", "description": "this param is a boolean", "authServices": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewBooleanParameterWithAuth("my_bool", "this param is a boolean", authServices), }, }, { name: "bool with authSources", in: []map[string]any{ { "name": "my_bool", "type": "boolean", "description": "this param is a boolean", "authSources": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewBooleanParameterWithAuth("my_bool", "this param is a boolean", authServices), }, }, { name: "string array", in: []map[string]any{ { "name": "my_array", "type": "array", "description": "this param is an array of strings", "items": map[string]string{ "name": "my_string", "type": "string", "description": "string item", }, "authServices": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewArrayParameterWithAuth("my_array", "this param is an array of strings", tools.NewStringParameter("my_string", "string item"), authServices), }, }, { name: "string array with authSources", in: []map[string]any{ { "name": "my_array", "type": "array", "description": "this param is an array of strings", "items": map[string]string{ "name": "my_string", "type": "string", "description": "string item", }, "authSources": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewArrayParameterWithAuth("my_array", "this param is an array of strings", tools.NewStringParameter("my_string", "string item"), authServices), }, }, { name: "float array", in: []map[string]any{ { "name": "my_array", "type": "array", "description": "this param is an array of floats", "items": map[string]string{ "name": "my_float", "type": "float", "description": "float item", }, "authServices": []map[string]string{ { "name": "my-google-auth-service", "field": "user_id", }, { "name": "other-auth-service", "field": "user_id", }, }, }, }, want: tools.Parameters{ tools.NewArrayParameterWithAuth("my_array", "this param is an array of floats", tools.NewFloatParameter("my_float", "float item"), authServices), }, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { var got tools.Parameters // parse map to bytes data, err := yaml.Marshal(tc.in) if err != nil { t.Fatalf("unable to marshal input to yaml: %s", err) } // parse bytes to object err = yaml.UnmarshalContext(ctx, data, &got) if err != nil { t.Fatalf("unable to unmarshal: %s", err) } if diff := cmp.Diff(tc.want, got); diff != "" { t.Fatalf("incorrect parse: diff %v", diff) } }) } } func TestParametersParse(t *testing.T) { tcs := []struct { name string params tools.Parameters in map[string]any want tools.ParamValues }{ { name: "string", params: tools.Parameters{ tools.NewStringParameter("my_string", "this param is a string"), }, in: map[string]any{ "my_string": "hello world", }, want: tools.ParamValues{tools.ParamValue{Name: "my_string", Value: "hello world"}}, }, { name: "not string", params: tools.Parameters{ tools.NewStringParameter("my_string", "this param is a string"), }, in: map[string]any{ "my_string": 4, }, }, { name: "int", params: tools.Parameters{ tools.NewIntParameter("my_int", "this param is an int"), }, in: map[string]any{ "my_int": 100, }, want: tools.ParamValues{tools.ParamValue{Name: "my_int", Value: 100}}, }, { name: "not int", params: tools.Parameters{ tools.NewIntParameter("my_int", "this param is an int"), }, in: map[string]any{ "my_int": 14.5, }, }, { name: "not int (big)", params: tools.Parameters{ tools.NewIntParameter("my_int", "this param is an int"), }, in: map[string]any{ "my_int": math.MaxInt64, }, want: tools.ParamValues{tools.ParamValue{Name: "my_int", Value: math.MaxInt64}}, }, { name: "float", params: tools.Parameters{ tools.NewFloatParameter("my_float", "this param is a float"), }, in: map[string]any{ "my_float": 1.5, }, want: tools.ParamValues{tools.ParamValue{Name: "my_float", Value: 1.5}}, }, { name: "not float", params: tools.Parameters{ tools.NewFloatParameter("my_float", "this param is a float"), }, in: map[string]any{ "my_float": true, }, }, { name: "bool", params: tools.Parameters{ tools.NewBooleanParameter("my_bool", "this param is a bool"), }, in: map[string]any{ "my_bool": true, }, want: tools.ParamValues{tools.ParamValue{Name: "my_bool", Value: true}}, }, { name: "not bool", params: tools.Parameters{ tools.NewBooleanParameter("my_bool", "this param is a bool"), }, in: map[string]any{ "my_bool": 1.5, }, }, { name: "string default", params: tools.Parameters{ tools.NewStringParameterWithDefault("my_string", "foo", "this param is a string"), }, in: map[string]any{}, want: tools.ParamValues{tools.ParamValue{Name: "my_string", Value: "foo"}}, }, { name: "int default", params: tools.Parameters{ tools.NewIntParameterWithDefault("my_int", 100, "this param is an int"), }, in: map[string]any{}, want: tools.ParamValues{tools.ParamValue{Name: "my_int", Value: 100}}, }, { name: "int (big)", params: tools.Parameters{ tools.NewIntParameterWithDefault("my_big_int", math.MaxInt64, "this param is an int"), }, in: map[string]any{}, want: tools.ParamValues{tools.ParamValue{Name: "my_big_int", Value: math.MaxInt64}}, }, { name: "float default", params: tools.Parameters{ tools.NewFloatParameterWithDefault("my_float", 1.1, "this param is a float"), }, in: map[string]any{}, want: tools.ParamValues{tools.ParamValue{Name: "my_float", Value: 1.1}}, }, { name: "bool default", params: tools.Parameters{ tools.NewBooleanParameterWithDefault("my_bool", true, "this param is a bool"), }, in: map[string]any{}, want: tools.ParamValues{tools.ParamValue{Name: "my_bool", Value: true}}, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { // parse map to bytes data, err := json.Marshal(tc.in) if err != nil { t.Fatalf("unable to marshal input to yaml: %s", err) } // parse bytes to object var m map[string]any d := json.NewDecoder(bytes.NewReader(data)) d.UseNumber() err = d.Decode(&m) if err != nil { t.Fatalf("unable to unmarshal: %s", err) } wantErr := len(tc.want) == 0 // error is expected if no items in want gotAll, err := tools.ParseParams(tc.params, m, make(map[string]map[string]any)) if err != nil { if wantErr { return } t.Fatalf("unexpected error from ParseParams: %s", err) } if wantErr { t.Fatalf("expected error but Param parsed successfully: %s", gotAll) } for i, got := range gotAll { want := tc.want[i] if got != want { t.Fatalf("unexpected value: got %q, want %q", got, want) } gotType, wantType := reflect.TypeOf(got), reflect.TypeOf(want) if gotType != wantType { t.Fatalf("unexpected value: got %q, want %q", got, want) } } }) } } func TestAuthParametersParse(t *testing.T) { authServices := []tools.ParamAuthService{ { Name: "my-google-auth-service", Field: "auth_field", }, { Name: "other-auth-service", Field: "other_auth_field", }} tcs := []struct { name string params tools.Parameters in map[string]any claimsMap map[string]map[string]any want tools.ParamValues }{ { name: "string", params: tools.Parameters{ tools.NewStringParameterWithAuth("my_string", "this param is a string", authServices), }, in: map[string]any{ "my_string": "hello world", }, claimsMap: map[string]map[string]any{"my-google-auth-service": {"auth_field": "hello"}}, want: tools.ParamValues{tools.ParamValue{Name: "my_string", Value: "hello"}}, }, { name: "not string", params: tools.Parameters{ tools.NewStringParameterWithAuth("my_string", "this param is a string", authServices), }, in: map[string]any{ "my_string": 4, }, claimsMap: map[string]map[string]any{}, }, { name: "int", params: tools.Parameters{ tools.NewIntParameterWithAuth("my_int", "this param is an int", authServices), }, in: map[string]any{ "my_int": 100, }, claimsMap: map[string]map[string]any{"other-auth-service": {"other_auth_field": 120}}, want: tools.ParamValues{tools.ParamValue{Name: "my_int", Value: 120}}, }, { name: "not int", params: tools.Parameters{ tools.NewIntParameterWithAuth("my_int", "this param is an int", authServices), }, in: map[string]any{ "my_int": 14.5, }, claimsMap: map[string]map[string]any{}, }, { name: "float", params: tools.Parameters{ tools.NewFloatParameterWithAuth("my_float", "this param is a float", authServices), }, in: map[string]any{ "my_float": 1.5, }, claimsMap: map[string]map[string]any{"my-google-auth-service": {"auth_field": 2.1}}, want: tools.ParamValues{tools.ParamValue{Name: "my_float", Value: 2.1}}, }, { name: "not float", params: tools.Parameters{ tools.NewFloatParameterWithAuth("my_float", "this param is a float", authServices), }, in: map[string]any{ "my_float": true, }, claimsMap: map[string]map[string]any{}, }, { name: "bool", params: tools.Parameters{ tools.NewBooleanParameterWithAuth("my_bool", "this param is a bool", authServices), }, in: map[string]any{ "my_bool": true, }, claimsMap: map[string]map[string]any{"my-google-auth-service": {"auth_field": false}}, want: tools.ParamValues{tools.ParamValue{Name: "my_bool", Value: false}}, }, { name: "not bool", params: tools.Parameters{ tools.NewBooleanParameterWithAuth("my_bool", "this param is a bool", authServices), }, in: map[string]any{ "my_bool": 1.5, }, claimsMap: map[string]map[string]any{}, }, { name: "username", params: tools.Parameters{ tools.NewStringParameterWithAuth("username", "username string", authServices), }, in: map[string]any{ "username": "Violet", }, claimsMap: map[string]map[string]any{"my-google-auth-service": {"auth_field": "Alice"}}, want: tools.ParamValues{tools.ParamValue{Name: "username", Value: "Alice"}}, }, { name: "expect claim error", params: tools.Parameters{ tools.NewStringParameterWithAuth("username", "username string", authServices), }, in: map[string]any{ "username": "Violet", }, claimsMap: map[string]map[string]any{"my-google-auth-service": {"not_an_auth_field": "Alice"}}, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { // parse map to bytes data, err := json.Marshal(tc.in) if err != nil { t.Fatalf("unable to marshal input to yaml: %s", err) } // parse bytes to object var m map[string]any d := json.NewDecoder(bytes.NewReader(data)) d.UseNumber() err = d.Decode(&m) if err != nil { t.Fatalf("unable to unmarshal: %s", err) } gotAll, err := tools.ParseParams(tc.params, m, tc.claimsMap) if err != nil { if len(tc.want) == 0 { // error is expected if no items in want return } t.Fatalf("unexpected error from ParseParams: %s", err) } for i, got := range gotAll { want := tc.want[i] if got != want { t.Fatalf("unexpected value: got %q, want %q", got, want) } gotType, wantType := reflect.TypeOf(got), reflect.TypeOf(want) if gotType != wantType { t.Fatalf("unexpected value: got %q, want %q", got, want) } } }) } } func TestParamValues(t *testing.T) { tcs := []struct { name string in tools.ParamValues wantSlice []any wantMap map[string]interface{} wantMapOrdered map[string]interface{} wantMapWithDollar map[string]interface{} }{ { name: "string", in: tools.ParamValues{tools.ParamValue{Name: "my_bool", Value: true}, tools.ParamValue{Name: "my_string", Value: "hello world"}}, wantSlice: []any{true, "hello world"}, wantMap: map[string]interface{}{"my_bool": true, "my_string": "hello world"}, wantMapOrdered: map[string]interface{}{"p1": true, "p2": "hello world"}, wantMapWithDollar: map[string]interface{}{ "$my_bool": true, "$my_string": "hello world", }, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { gotSlice := tc.in.AsSlice() gotMap := tc.in.AsMap() gotMapOrdered := tc.in.AsMapByOrderedKeys() gotMapWithDollar := tc.in.AsMapWithDollarPrefix() for i, got := range gotSlice { want := tc.wantSlice[i] if got != want { t.Fatalf("unexpected value: got %q, want %q", got, want) } } for i, got := range gotMap { want := tc.wantMap[i] if got != want { t.Fatalf("unexpected value: got %q, want %q", got, want) } } for i, got := range gotMapOrdered { want := tc.wantMapOrdered[i] if got != want { t.Fatalf("unexpected value: got %q, want %q", got, want) } } for key, got := range gotMapWithDollar { want := tc.wantMapWithDollar[key] if got != want { t.Fatalf("unexpected value in AsMapWithDollarPrefix: got %q, want %q", got, want) } } }) } } func TestParamManifest(t *testing.T) { tcs := []struct { name string in tools.Parameter want tools.ParameterManifest }{ { name: "string", in: tools.NewStringParameter("foo-string", "bar"), want: tools.ParameterManifest{Name: "foo-string", Type: "string", Required: true, Description: "bar", AuthServices: []string{}}, }, { name: "int", in: tools.NewIntParameter("foo-int", "bar"), want: tools.ParameterManifest{Name: "foo-int", Type: "integer", Required: true, Description: "bar", AuthServices: []string{}}, }, { name: "float", in: tools.NewFloatParameter("foo-float", "bar"), want: tools.ParameterManifest{Name: "foo-float", Type: "float", Required: true, Description: "bar", AuthServices: []string{}}, }, { name: "boolean", in: tools.NewBooleanParameter("foo-bool", "bar"), want: tools.ParameterManifest{Name: "foo-bool", Type: "boolean", Required: true, Description: "bar", AuthServices: []string{}}, }, { name: "array", in: tools.NewArrayParameter("foo-array", "bar", tools.NewStringParameter("foo-string", "bar")), want: tools.ParameterManifest{ Name: "foo-array", Type: "array", Required: true, Description: "bar", AuthServices: []string{}, Items: &tools.ParameterManifest{Name: "foo-string", Type: "string", Required: true, Description: "bar", AuthServices: []string{}}, }, }, { name: "string default", in: tools.NewStringParameterWithDefault("foo-string", "foo", "bar"), want: tools.ParameterManifest{Name: "foo-string", Type: "string", Required: false, Description: "bar", AuthServices: []string{}}, }, { name: "int default", in: tools.NewIntParameterWithDefault("foo-int", 1, "bar"), want: tools.ParameterManifest{Name: "foo-int", Type: "integer", Required: false, Description: "bar", AuthServices: []string{}}, }, { name: "float default", in: tools.NewFloatParameterWithDefault("foo-float", 1.1, "bar"), want: tools.ParameterManifest{Name: "foo-float", Type: "float", Required: false, Description: "bar", AuthServices: []string{}}, }, { name: "boolean default", in: tools.NewBooleanParameterWithDefault("foo-bool", true, "bar"), want: tools.ParameterManifest{Name: "foo-bool", Type: "boolean", Required: false, Description: "bar", AuthServices: []string{}}, }, { name: "array default", in: tools.NewArrayParameterWithDefault("foo-array", `["foo", "bar"]`, "bar", tools.NewStringParameter("foo-string", "bar")), want: tools.ParameterManifest{ Name: "foo-array", Type: "array", Required: false, Description: "bar", AuthServices: []string{}, Items: &tools.ParameterManifest{Name: "foo-string", Type: "string", Required: true, Description: "bar", AuthServices: []string{}}, }, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { got := tc.in.Manifest() if !reflect.DeepEqual(got, tc.want) { t.Fatalf("unexpected manifest: got %+v, want %+v", got, tc.want) } }) } } func TestParamMcpManifest(t *testing.T) { tcs := []struct { name string in tools.Parameter want tools.ParameterMcpManifest }{ { name: "string", in: tools.NewStringParameter("foo-string", "bar"), want: tools.ParameterMcpManifest{Type: "string", Description: "bar"}, }, { name: "int", in: tools.NewIntParameter("foo-int", "bar"), want: tools.ParameterMcpManifest{Type: "integer", Description: "bar"}, }, { name: "float", in: tools.NewFloatParameter("foo-float", "bar"), want: tools.ParameterMcpManifest{Type: "float", Description: "bar"}, }, { name: "boolean", in: tools.NewBooleanParameter("foo-bool", "bar"), want: tools.ParameterMcpManifest{Type: "boolean", Description: "bar"}, }, { name: "array", in: tools.NewArrayParameter("foo-array", "bar", tools.NewStringParameter("foo-string", "bar")), want: tools.ParameterMcpManifest{ Type: "array", Description: "bar", Items: &tools.ParameterMcpManifest{Type: "string", Description: "bar"}, }, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { got := tc.in.McpManifest() if !reflect.DeepEqual(got, tc.want) { t.Fatalf("unexpected manifest: got %+v, want %+v", got, tc.want) } }) } } func TestMcpManifest(t *testing.T) { tcs := []struct { name string in tools.Parameters want tools.McpToolsSchema }{ { name: "string", in: tools.Parameters{ tools.NewStringParameterWithDefault("foo-string", "foo", "bar"), tools.NewStringParameter("foo-string2", "bar"), tools.NewIntParameterWithDefault("foo-int", 1, "bar"), tools.NewIntParameter("foo-int2", "bar"), tools.NewArrayParameterWithDefault("foo-array", []string{"hello", "world"}, "bar", tools.NewStringParameter("foo-string", "bar")), tools.NewArrayParameter("foo-array2", "bar", tools.NewStringParameter("foo-string", "bar")), }, want: tools.McpToolsSchema{ Type: "object", Properties: map[string]tools.ParameterMcpManifest{ "foo-string": tools.ParameterMcpManifest{Type: "string", Description: "bar"}, "foo-string2": tools.ParameterMcpManifest{Type: "string", Description: "bar"}, "foo-int": tools.ParameterMcpManifest{Type: "integer", Description: "bar"}, "foo-int2": tools.ParameterMcpManifest{Type: "integer", Description: "bar"}, "foo-array": tools.ParameterMcpManifest{ Type: "array", Description: "bar", Items: &tools.ParameterMcpManifest{Type: "string", Description: "bar"}, }, "foo-array2": tools.ParameterMcpManifest{ Type: "array", Description: "bar", Items: &tools.ParameterMcpManifest{Type: "string", Description: "bar"}, }, }, Required: []string{"foo-string2", "foo-int2", "foo-array2"}, }, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { got := tc.in.McpManifest() if !reflect.DeepEqual(got, tc.want) { t.Fatalf("unexpected manifest: got %+v, want %+v", got, tc.want) } }) } } func TestFailParametersUnmarshal(t *testing.T) { ctx, err := testutils.ContextWithNewLogger() if err != nil { t.Fatalf("unexpected error: %s", err) } tcs := []struct { name string in []map[string]any err string }{ { name: "common parameter missing name", in: []map[string]any{ { "type": "string", "description": "this is a param for string", }, }, err: "unable to parse as \"string\": Key: 'CommonParameter.Name' Error:Field validation for 'Name' failed on the 'required' tag", }, { name: "common parameter missing type", in: []map[string]any{ { "name": "string", "description": "this is a param for string", }, }, err: "parameter is missing 'type' field: %!w()", }, { name: "common parameter missing description", in: []map[string]any{ { "name": "my_string", "type": "string", }, }, err: "unable to parse as \"string\": Key: 'CommonParameter.Desc' Error:Field validation for 'Desc' failed on the 'required' tag", }, { name: "array parameter missing items", in: []map[string]any{ { "name": "my_array", "type": "array", "description": "this param is an array of strings", }, }, err: "unable to parse as \"array\": unable to parse 'items' field: error parsing parameters: nothing to unmarshal", }, { name: "array parameter missing items' name", in: []map[string]any{ { "name": "my_array", "type": "array", "description": "this param is an array of strings", "items": map[string]string{ "type": "string", "description": "string item", }, }, }, err: "unable to parse as \"array\": unable to parse 'items' field: unable to parse as \"string\": Key: 'CommonParameter.Name' Error:Field validation for 'Name' failed on the 'required' tag", }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { var got tools.Parameters // parse map to bytes data, err := yaml.Marshal(tc.in) if err != nil { t.Fatalf("unable to marshal input to yaml: %s", err) } // parse bytes to object err = yaml.UnmarshalContext(ctx, data, &got) if err == nil { t.Fatalf("expect parsing to fail") } errStr := err.Error() if errStr != tc.err { t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err) } }) } } func TestConvertArrayParamToString(t *testing.T) { tcs := []struct { name string in []any want string }{ { in: []any{ "id", "name", "location", }, want: "id, name, location", }, { in: []any{ "id", }, want: "id", }, { in: []any{ "id", "5", "false", }, want: "id, 5, false", }, { in: []any{}, want: "", }, { in: []any{}, want: "", }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { got, _ := tools.ConvertArrayParamToString(tc.in) if diff := cmp.Diff(tc.want, got); diff != "" { t.Fatalf("incorrect array param conversion: diff %v", diff) } }) } } func TestFailConvertArrayParamToString(t *testing.T) { tcs := []struct { name string in []any err string }{ { in: []any{5, 10, 15}, err: "templateParameter only supports string arrays", }, { in: []any{"id", "name", 15}, err: "templateParameter only supports string arrays", }, { in: []any{false}, err: "templateParameter only supports string arrays", }, { in: []any{10, true}, err: "templateParameter only supports string arrays", }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { _, err := tools.ConvertArrayParamToString(tc.in) errStr := err.Error() if errStr != tc.err { t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err) } }) } } func TestGetParams(t *testing.T) { tcs := []struct { name string in map[string]any params tools.Parameters want tools.ParamValues }{ { name: "parameters to include and exclude", params: tools.Parameters{ tools.NewStringParameter("my_string_inc", "this should be included"), tools.NewStringParameter("my_string_inc2", "this should be included"), }, in: map[string]any{ "my_string_inc": "hello world A", "my_string_inc2": "hello world B", "my_string_exc": "hello world C", }, want: tools.ParamValues{ tools.ParamValue{Name: "my_string_inc", Value: "hello world A"}, tools.ParamValue{Name: "my_string_inc2", Value: "hello world B"}, }, }, { name: "include all", params: tools.Parameters{ tools.NewStringParameter("my_string_inc", "this should be included"), }, in: map[string]any{ "my_string_inc": "hello world A", }, want: tools.ParamValues{ tools.ParamValue{Name: "my_string_inc", Value: "hello world A"}, }, }, { name: "exclude all", params: tools.Parameters{}, in: map[string]any{ "my_string_exc": "hello world A", "my_string_exc2": "hello world B", }, want: tools.ParamValues{}, }, { name: "empty", params: tools.Parameters{}, in: map[string]any{}, want: tools.ParamValues{}, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { got, _ := tools.GetParams(tc.params, tc.in) if diff := cmp.Diff(tc.want, got); diff != "" { t.Fatalf("incorrect get params: diff %v", diff) } }) } } func TestFailGetParams(t *testing.T) { tcs := []struct { name string params tools.Parameters in map[string]any err string }{ { name: "missing the only parameter", params: tools.Parameters{tools.NewStringParameter("my_string", "this was missing")}, in: map[string]any{}, err: "missing parameter my_string", }, { name: "missing one parameter of multiple", params: tools.Parameters{ tools.NewStringParameter("my_string_inc", "this should be included"), tools.NewStringParameter("my_string_exc", "this was missing"), }, in: map[string]any{ "my_string_inc": "hello world A", }, err: "missing parameter my_string_exc", }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { _, err := tools.GetParams(tc.params, tc.in) errStr := err.Error() if errStr != tc.err { t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err) } }) } } func TestResolveTemplateParameters(t *testing.T) { tcs := []struct { name string templateParams tools.Parameters statement string in map[string]any want string }{ { name: "single template parameter", templateParams: tools.Parameters{ tools.NewStringParameter("tableName", "this is a string template parameter"), }, statement: "SELECT * FROM {{.tableName}}", in: map[string]any{ "tableName": "hotels", }, want: "SELECT * FROM hotels", }, { name: "multiple template parameters", templateParams: tools.Parameters{ tools.NewStringParameter("tableName", "this is a string template parameter"), tools.NewStringParameter("columnName", "this is a string template parameter"), }, statement: "SELECT * FROM {{.tableName}} WHERE {{.columnName}} = 'Hilton'", in: map[string]any{ "tableName": "hotels", "columnName": "name", }, want: "SELECT * FROM hotels WHERE name = 'Hilton'", }, { name: "standard and template parameter", templateParams: tools.Parameters{ tools.NewStringParameter("tableName", "this is a string template parameter"), tools.NewStringParameter("hotelName", "this is a string parameter"), }, statement: "SELECT * FROM {{.tableName}} WHERE name = $1", in: map[string]any{ "tableName": "hotels", "hotelName": "name", }, want: "SELECT * FROM hotels WHERE name = $1", }, { name: "standard parameter", templateParams: tools.Parameters{ tools.NewStringParameter("hotelName", "this is a string parameter"), }, statement: "SELECT * FROM hotels WHERE name = $1", in: map[string]any{ "hotelName": "hotels", }, want: "SELECT * FROM hotels WHERE name = $1", }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { got, _ := tools.ResolveTemplateParams(tc.templateParams, tc.statement, tc.in) if diff := cmp.Diff(tc.want, got); diff != "" { t.Fatalf("incorrect resolved template params: diff %v", diff) } }) } } func TestFailResolveTemplateParameters(t *testing.T) { tcs := []struct { name string templateParams tools.Parameters statement string in map[string]any err string }{ { name: "wrong param name", templateParams: tools.Parameters{ tools.NewStringParameter("tableName", "this is a string template parameter"), }, statement: "SELECT * FROM {{.missingParam}}", in: map[string]any{}, err: "error getting template params missing parameter tableName", }, { name: "incomplete param template", templateParams: tools.Parameters{ tools.NewStringParameter("tableName", "this is a string template parameter"), }, statement: "SELECT * FROM {{.tableName", in: map[string]any{ "tableName": "hotels", }, err: "error creating go template template: statement:1: unclosed action", }, { name: "undefined function", templateParams: tools.Parameters{ tools.NewStringParameter("tableName", "this is a string template parameter"), }, statement: "SELECT * FROM {{json .tableName}}", in: map[string]any{ "tableName": "hotels", }, err: "error creating go template template: statement:1: function \"json\" not defined", }, { name: "undefined method", templateParams: tools.Parameters{ tools.NewStringParameter("tableName", "this is a string template parameter"), }, statement: "SELECT * FROM {{.tableName .wrong}}", in: map[string]any{ "tableName": "hotels", }, err: "error executing go template template: statement:1:16: executing \"statement\" at <.tableName>: tableName is not a method but has arguments", }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { _, err := tools.ResolveTemplateParams(tc.templateParams, tc.statement, tc.in) errStr := err.Error() if errStr != tc.err { t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err) } }) } }