genai-toolbox/internal/prompts/arguments_test.go

// 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 (
	"fmt"
	"strings"
	"testing"

	yaml "github.com/goccy/go-yaml"
	"github.com/google/go-cmp/cmp"
	"github.com/googleapis/genai-toolbox/internal/prompts"
	"github.com/googleapis/genai-toolbox/internal/testutils"
	"github.com/googleapis/genai-toolbox/internal/util/parameters"
)

// Test type aliases for convenience.
type (
	Argument       = prompts.Argument
	ArgMcpManifest = prompts.ArgMcpManifest
	Arguments      = prompts.Arguments
)

// Ptr is a helper function to create a pointer to a value.
func Ptr[T any](v T) *T {
	return &v
}

func makeArrayArg(name, desc string, items parameters.Parameter) Argument {
	return Argument{Parameter: parameters.NewArrayParameter(name, desc, items)}
}

func TestArgMcpManifest(t *testing.T) {
	t.Parallel()
	testCases := []struct {
		name     string
		arg      Argument
		expected ArgMcpManifest
	}{
		{
			name: "Required with no default",
			arg:  Argument{Parameter: parameters.NewStringParameterWithRequired("name1", "desc1", true)},
			expected: ArgMcpManifest{
				Name: "name1", Description: "desc1", Required: true,
			},
		},
		{
			name: "Not required with no default",
			arg:  Argument{Parameter: parameters.NewStringParameterWithRequired("name2", "desc2", false)},
			expected: ArgMcpManifest{
				Name: "name2", Description: "desc2", Required: false,
			},
		},
		{
			name: "Implicitly required with default",
			arg:  Argument{Parameter: parameters.NewStringParameterWithDefault("name3", "defaultVal", "desc3")},
			expected: ArgMcpManifest{
				Name: "name3", Description: "desc3", Required: false,
			},
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			got := tc.arg.McpManifest()
			if diff := cmp.Diff(tc.expected, got); diff != "" {
				t.Errorf("McpManifest() mismatch (-want +got):\n%s", diff)
			}
		})
	}
}

// TestArguments_UnmarshalYAML tests all unmarshaling logic for the Arguments type.
func TestArgumentsUnmarshalYAML(t *testing.T) {
	t.Parallel()
	// paramComparer allows cmp.Diff to intelligently compare the parsed results.
	var transformFunc func(parameters.Parameter) any
	transformFunc = func(p parameters.Parameter) any {
		s := struct{ Name, Type, Desc string }{
			Name: p.GetName(),
			Type: p.GetType(),
			Desc: p.Manifest().Description,
		}
		if arr, ok := p.(*parameters.ArrayParameter); ok {
			s.Desc = fmt.Sprintf("%s items:%v", s.Desc, transformFunc(arr.GetItems()))
		}
		return s
	}
	paramComparer := cmp.Transformer("Parameter", transformFunc)

	testCases := []struct {
		name         string
		yamlInput    []map[string]any
		expectedArgs Arguments
		wantErr      string
	}{
		{
			name: "Defaults type to string when omitted",
			yamlInput: []map[string]any{
				{"name": "p1", "description": "d1"},
			},
			expectedArgs: Arguments{
				{Parameter: parameters.NewStringParameter("p1", "d1")},
			},
		},
		{
			name: "Respects type when present",
			yamlInput: []map[string]any{
				{"name": "p1", "description": "d1", "type": "integer"},
			},
			expectedArgs: Arguments{
				{Parameter: parameters.NewIntParameter("p1", "d1")},
			},
		},
		{
			name: "Parses complex types like arrays correctly",
			yamlInput: []map[string]any{
				{
					"name":        "param_array",
					"description": "an array",
					"type":        "array",
					"items": map[string]any{
						"name":        "item_name",
						"type":        "string",
						"description": "an item",
					},
				},
			},
			expectedArgs: Arguments{
				makeArrayArg("param_array", "an array", parameters.NewStringParameter("item_name", "an item")),
			},
		},
		{
			name: "Propagates parsing error for unsupported type",
			yamlInput: []map[string]any{
				{"name": "p1", "description": "d1", "type": "unsupported"},
			},
			wantErr: `"unsupported" is not valid type for a parameter`,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			yamlBytes, err := yaml.Marshal(tc.yamlInput)
			if err != nil {
				t.Fatalf("Test setup failure: could not marshal test input to YAML: %v", err)
			}
			var got Arguments
			ctx, err := testutils.ContextWithNewLogger()
			if err != nil {
				t.Fatalf("Failed to create logger using testutils: %v", err)
			}
			err = yaml.UnmarshalContext(ctx, yamlBytes, &got)

			if tc.wantErr != "" {
				if err == nil {
					t.Fatalf("UnmarshalContext() expected error but got nil")
				}
				if !strings.Contains(err.Error(), tc.wantErr) {
					t.Errorf("UnmarshalContext() error mismatch:\nwant to contain: %q\ngot: %q", tc.wantErr, err.Error())
				}
			} else {
				if err != nil {
					t.Fatalf("UnmarshalContext() returned unexpected error: %v", err)
				}
				if diff := cmp.Diff(tc.expectedArgs, got, paramComparer); diff != "" {
					t.Errorf("UnmarshalContext() result mismatch (-want +got):\n%s", diff)
				}
			}
		})
	}
}

func TestParseArguments(t *testing.T) {
	t.Parallel()
	testArguments := prompts.Arguments{
		{Parameter: parameters.NewStringParameter("name", "A required name.")},
		{Parameter: parameters.NewIntParameterWithRequired("count", "An optional count.", false)},
	}

	testCases := []struct {
		name    string
		argsIn  map[string]any
		want    parameters.ParamValues
		wantErr string
	}{
		{
			name: "Success with all parameters provided",
			argsIn: map[string]any{
				"name":  "test-name",
				"count": 42,
			},
			want: parameters.ParamValues{
				{Name: "name", Value: "test-name"},
				{Name: "count", Value: 42},
			},
		},
		{
			name: "Success with only required parameters",
			argsIn: map[string]any{
				"name": "another-name",
			},
			want: parameters.ParamValues{
				{Name: "name", Value: "another-name"},
				{Name: "count", Value: nil},
			},
		},
		{
			name: "Failure with missing required parameter",
			argsIn: map[string]any{
				"count": 123,
			},
			wantErr: `parameter "name" is required`,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			got, err := prompts.ParseArguments(testArguments, tc.argsIn, nil)
			if tc.wantErr != "" {
				if err == nil {
					t.Fatalf("expected an error but got nil")
				}
				if !strings.Contains(err.Error(), tc.wantErr) {
					t.Errorf("error mismatch:\n  want to contain: %q\n  got: %q", tc.wantErr, err.Error())
				}
			} else {
				if err != nil {
					t.Fatalf("unexpected error: %v", err)
				}
				if diff := cmp.Diff(tc.want, got); diff != "" {
					t.Errorf("ParseArguments() result mismatch (-want +got):\n%s", diff)
				}
			}
		})
	}
}