MCP Toolbox for Databases

// Copyright 2025 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 prompts_test import ( "strings" "testing" "github.com/google/go-cmp/cmp" "github.com/googleapis/genai-toolbox/internal/prompts" "github.com/googleapis/genai-toolbox/internal/util/parameters" ) // mockPrompt is a simple mock implementation of prompts.Prompt for testing. type mockPrompt struct { name string desc string args prompts.Arguments manifest prompts.Manifest mcpManifest prompts.McpManifest } func (m mockPrompt) SubstituteParams(parameters.ParamValues) (any, error) { return nil, nil } func (m mockPrompt) ParseArgs(map[string]any, map[string]map[string]any) (parameters.ParamValues, error) { return nil, nil } func (m mockPrompt) Manifest() prompts.Manifest { return m.manifest } func (m mockPrompt) McpManifest() prompts.McpManifest { return m.mcpManifest } func (m mockPrompt) ToConfig() prompts.PromptConfig { return nil } // newMockPrompt creates a new mock prompt for testing. func newMockPrompt(name, desc string) prompts.Prompt { args := prompts.Arguments{ {Parameter: parameters.NewStringParameter("arg1", "Test argument")}, } return mockPrompt{ name: name, desc: desc, args: args, manifest: prompts.Manifest{ Description: desc, Arguments: []parameters.ParameterManifest{ {Name: "arg1", Type: "string", Required: true, Description: "Test argument", AuthServices: []string{}}, }, }, mcpManifest: prompts.McpManifest{ Name: name, Description: desc, Arguments: []prompts.ArgMcpManifest{ {Name: "arg1", Description: "Test argument", Required: true}, }, }, } } func TestPromptsetConfig_Initialize(t *testing.T) { t.Parallel() promptsMap := map[string]prompts.Prompt{ "prompt1": newMockPrompt("prompt1", "First test prompt"), "prompt2": newMockPrompt("prompt2", "Second test prompt"), } serverVersion := "v1.0.0" p1 := promptsMap["prompt1"] p2 := promptsMap["prompt2"] prompt1Ptr := &p1 prompt2Ptr := &p2 testCases := []struct { name string config prompts.PromptsetConfig want prompts.Promptset wantErr string }{ { name: "Success case", config: prompts.PromptsetConfig{ Name: "default", PromptNames: []string{"prompt1", "prompt2"}, }, want: prompts.Promptset{ PromptsetConfig: prompts.PromptsetConfig{ Name: "default", }, Prompts: []*prompts.Prompt{ prompt1Ptr, prompt2Ptr, }, Manifest: prompts.PromptsetManifest{ ServerVersion: serverVersion, PromptsManifest: map[string]prompts.Manifest{ "prompt1": promptsMap["prompt1"].Manifest(), "prompt2": promptsMap["prompt2"].Manifest(), }, }, McpManifest: []prompts.McpManifest{ promptsMap["prompt1"].McpManifest(), promptsMap["prompt2"].McpManifest(), }, }, wantErr: "", }, { name: "Success case with one prompt", config: prompts.PromptsetConfig{ Name: "single", PromptNames: []string{"prompt1"}, }, want: prompts.Promptset{ PromptsetConfig: prompts.PromptsetConfig{ Name: "single", }, Prompts: []*prompts.Prompt{ prompt1Ptr, }, Manifest: prompts.PromptsetManifest{ ServerVersion: serverVersion, PromptsManifest: map[string]prompts.Manifest{ "prompt1": promptsMap["prompt1"].Manifest(), }, }, McpManifest: []prompts.McpManifest{ promptsMap["prompt1"].McpManifest(), }, }, wantErr: "", }, { name: "Failure case - invalid promptset name", config: prompts.PromptsetConfig{ Name: "invalid name", // Contains a space PromptNames: []string{"prompt1"}, }, want: prompts.Promptset{PromptsetConfig: prompts.PromptsetConfig{Name: "invalid name"}}, // Expect partial struct wantErr: "invalid promptset name", }, { name: "Failure case - prompt not found", config: prompts.PromptsetConfig{ Name: "missing_prompt", PromptNames: []string{"prompt1", "prompt_does_not_exist"}, }, // Expect partial struct with fields populated up to the error want: prompts.Promptset{ PromptsetConfig: prompts.PromptsetConfig{ Name: "missing_prompt", }, Prompts: []*prompts.Prompt{ prompt1Ptr, }, Manifest: prompts.PromptsetManifest{ ServerVersion: serverVersion, PromptsManifest: map[string]prompts.Manifest{ "prompt1": promptsMap["prompt1"].Manifest(), }, }, McpManifest: []prompts.McpManifest{ promptsMap["prompt1"].McpManifest(), }, }, wantErr: "prompt does not exist", }, { name: "Success case - empty prompt list", config: prompts.PromptsetConfig{ Name: "empty", PromptNames: []string{}, }, want: prompts.Promptset{ PromptsetConfig: prompts.PromptsetConfig{ Name: "empty", }, Prompts: []*prompts.Prompt{}, Manifest: prompts.PromptsetManifest{ ServerVersion: serverVersion, PromptsManifest: map[string]prompts.Manifest{}, }, McpManifest: []prompts.McpManifest{}, }, wantErr: "", }, } for _, tc := range testCases { tc := tc t.Run(tc.name, func(t *testing.T) { got, err := tc.config.Initialize(serverVersion, promptsMap) if tc.wantErr != "" { if err == nil { t.Fatalf("Initialize() expected error but got nil") } if !strings.Contains(err.Error(), tc.wantErr) { t.Errorf("Initialize() error mismatch:\n want to contain: %q\n got: %q", tc.wantErr, err.Error()) } // Also check that the partially populated struct matches if diff := cmp.Diff(tc.want, got, cmp.AllowUnexported(mockPrompt{})); diff != "" { t.Errorf("Initialize() partial result on error mismatch (-want +got):\n%s", diff) } } else { if err != nil { t.Fatalf("Initialize() returned unexpected error: %v", err) } // Using cmp.AllowUnexported because mockPrompt is unexported if diff := cmp.Diff(tc.want, got, cmp.AllowUnexported(mockPrompt{})); diff != "" { t.Errorf("Initialize() result mismatch (-want +got):\n%s", diff) } } }) } }