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 custom import ( "bytes" "context" "encoding/json" "io" "net/http" "regexp" "testing" "time" _ "github.com/googleapis/genai-toolbox/internal/prompts/custom" "github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc" "github.com/googleapis/genai-toolbox/internal/testutils" "github.com/googleapis/genai-toolbox/tests" ) func getPromptsConfig() map[string]any { return map[string]any{ "prompts": map[string]any{ "prompt1": map[string]any{ "description": "description1.", "arguments": []any{ map[string]any{ "name": "arg1", "description": "Description of arg1 for prompt1", }, }, "messages": []any{ map[string]any{ "content": "content arg1: {{.arg1}}", }, }, }, "prompt2": map[string]any{ "description": "description2", "arguments": []any{ map[string]any{ "name": "arg1", "description": "Description for arg1 for prompt2", }, map[string]any{ "name": "arg2", "description": "Description for arg2 for prompt2", }, }, "messages": []any{ map[string]any{ "role": "user", "content": "message1 args: {{.arg1}}, {{.arg2}}", }, map[string]any{ "role": "assistant", "content": "{{.arg1}}", }, }, }, }, } } // TestMCPPromptsIntegration is the main entrypoint for the prompts integration tests. func TestMCPPromptsIntegration(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) defer cancel() cmd, cleanup, err := tests.StartCmd(ctx, getPromptsConfig()) if err != nil { t.Fatalf("command initialization returned an error: %v", err) } defer cleanup() waitCtx, cancelWait := context.WithTimeout(ctx, 20*time.Second) defer cancelWait() out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out) if err != nil { t.Logf("toolbox command logs: \n%s", out) t.Fatalf("toolbox didn't start successfully: %v", err) } runMCPListPromptsTest(t) runMCPGetPromptTest(t) } func runMCPListPromptsTest(t *testing.T) { t.Run("MCP_prompts/list", func(t *testing.T) { api := "http://127.0.0.1:5000/mcp" reqBody := jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "list-prompts-1", Request: jsonrpc.Request{Method: "prompts/list"}, } jsonBody, _ := json.Marshal(reqBody) resp, err := http.Post(api, "application/json", bytes.NewBuffer(jsonBody)) if err != nil { t.Fatalf("failed to send request: %v", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { t.Fatalf("expected status OK; got %s", resp.Status) } var result struct { Result struct { Prompts []map[string]any `json:"prompts"` } `json:"result"` } if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { t.Fatalf("failed to decode response: %v", err) } if len(result.Result.Prompts) != 2 { t.Errorf("expected 2 prompts, got %d", len(result.Result.Prompts)) } }) } func runMCPGetPromptTest(t *testing.T) { t.Run("MCP_prompts/get", func(t *testing.T) { api := "http://127.0.0.1:5000/mcp" reqBody := jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "get-prompt-2", Request: jsonrpc.Request{Method: "prompts/get"}, Params: map[string]any{ "name": "prompt2", "arguments": map[string]any{ "arg1": "value1", "arg2": "value2", }, }, } jsonBody, _ := json.Marshal(reqBody) resp, err := http.Post(api, "application/json", bytes.NewBuffer(jsonBody)) if err != nil { t.Fatalf("failed to send request: %v", err) } defer resp.Body.Close() body, _ := io.ReadAll(resp.Body) if resp.StatusCode != http.StatusOK { t.Fatalf("expected status OK; got %s, body: %s", resp.Status, string(body)) } var result map[string]any if err := json.Unmarshal(body, &result); err != nil { t.Fatalf("failed to decode response: %v", err) } messages := result["result"].(map[string]any)["messages"].([]any) if len(messages) != 2 { t.Fatalf("expected 2 messages, got %d", len(messages)) } // Check first message msg1 := messages[0].(map[string]any) content1 := msg1["content"].(map[string]any)["text"].(string) expectedContent1 := "message1 args: value1, value2" if content1 != expectedContent1 { t.Errorf("unexpected content in message 1.\nGot: %q\nWant: %q", content1, expectedContent1) } // Check second message msg2 := messages[1].(map[string]any) content2 := msg2["content"].(map[string]any)["text"].(string) expectedContent2 := "value1" if content2 != expectedContent2 { t.Errorf("unexpected content in message 2.\nGot: %q\nWant: %q", content2, expectedContent2) } }) }