genai-toolbox/internal/tools/cloudgda/cloudgda_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 cloudgda_test

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"

	yaml "github.com/goccy/go-yaml"
	"github.com/google/go-cmp/cmp"
	"github.com/googleapis/genai-toolbox/internal/server"
	"github.com/googleapis/genai-toolbox/internal/server/resources"
	"github.com/googleapis/genai-toolbox/internal/sources"
	cloudgdasrc "github.com/googleapis/genai-toolbox/internal/sources/cloudgda"
	"github.com/googleapis/genai-toolbox/internal/testutils"
	"github.com/googleapis/genai-toolbox/internal/tools"
	cloudgdatool "github.com/googleapis/genai-toolbox/internal/tools/cloudgda"
	"github.com/googleapis/genai-toolbox/internal/util/parameters"
)

func TestParseFromYaml(t *testing.T) {
	t.Parallel()
	tcs := []struct {
		desc string
		in   string
		want server.ToolConfigs
	}{
		{
			desc: "basic example",
			in: `
			tools:
				my-gda-query-tool:
					kind: cloud-gemini-data-analytics-query
					source: gda-api-source
					description: Test Description
					location: us-central1
					context:
						datasourceReferences:
							spannerReference:
								databaseReference:
									projectId:  "cloud-db-nl2sql"
									region:     "us-central1"
									instanceId: "evalbench"
									databaseId: "financial"
									engine:     "GOOGLE_SQL"
								agentContextReference:
									contextSetId: "projects/cloud-db-nl2sql/locations/us-east1/contextSets/bdf_gsql_gemini_all_templates"
					generationOptions:
						generateQueryResult: true
			`,
			want: map[string]tools.ToolConfig{
				"my-gda-query-tool": cloudgdatool.Config{
					Name:         "my-gda-query-tool",
					Kind:         "cloud-gemini-data-analytics-query",
					Source:       "gda-api-source",
					Description:  "Test Description",
					Location:     "us-central1",
					AuthRequired: []string{},
					Context: &cloudgdatool.QueryDataContext{
						DatasourceReferences: &cloudgdatool.DatasourceReferences{
							SpannerReference: &cloudgdatool.SpannerReference{
								DatabaseReference: &cloudgdatool.SpannerDatabaseReference{
									ProjectID:  "cloud-db-nl2sql",
									Region:     "us-central1",
									InstanceID: "evalbench",
									DatabaseID: "financial",
									Engine:     cloudgdatool.SpannerEngineGoogleSQL,
								},
								AgentContextReference: &cloudgdatool.AgentContextReference{
									ContextSetID: "projects/cloud-db-nl2sql/locations/us-east1/contextSets/bdf_gsql_gemini_all_templates",
								},
							},
						},
					},
					GenerationOptions: &cloudgdatool.GenerationOptions{
						GenerateQueryResult: true,
					},
				},
			},
		},
	}
	for _, tc := range tcs {
		tc := tc
		t.Run(tc.desc, func(t *testing.T) {
			t.Parallel()
			got := struct {
				Tools server.ToolConfigs `yaml:"tools"`
			}{}
			// Parse contents
			err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
			if err != nil {
				t.Fatalf("unable to unmarshal: %s", err)
			}
			if !cmp.Equal(tc.want, got.Tools) {
				t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Tools)
			}
		})
	}
}

// authRoundTripper is a mock http.RoundTripper that adds a dummy Authorization header.
type authRoundTripper struct {
	Token string
	Next  http.RoundTripper
}

func (rt *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	newReq := *req
	newReq.Header = make(http.Header)
	for k, v := range req.Header {
		newReq.Header[k] = v
	}
	newReq.Header.Set("Authorization", rt.Token)
	if rt.Next == nil {
		return http.DefaultTransport.RoundTrip(&newReq)
	}
	return rt.Next.RoundTrip(&newReq)
}

type mockSource struct {
	kind    string
	client  *http.Client       // Can be used to inject a specific client
	baseURL string             // BaseURL is needed to implement sources.Source.BaseURL
	config  cloudgdasrc.Config // to return from ToConfig
}

func (m *mockSource) SourceKind() string             { return m.kind }
func (m *mockSource) ToConfig() sources.SourceConfig { return m.config }
func (m *mockSource) GetClient(ctx context.Context, token string) (*http.Client, error) {
	if m.client != nil {
		return m.client, nil
	}
	// Default client for testing if not explicitly set
	transport := &http.Transport{}
	authTransport := &authRoundTripper{
		Token: "Bearer test-access-token", // Dummy token
		Next:  transport,
	}
	return &http.Client{Transport: authTransport}, nil
}
func (m *mockSource) UseClientAuthorization() bool { return false }
func (m *mockSource) Initialize(ctx context.Context, tracer interface{}) (sources.Source, error) {
	return m, nil
}
func (m *mockSource) BaseURL() string { return m.baseURL }

func TestInitialize(t *testing.T) {
	t.Parallel()

	srcs := map[string]sources.Source{
		"gda-api-source": &cloudgdasrc.Source{
			Config:  cloudgdasrc.Config{Name: "gda-api-source", Kind: cloudgdasrc.SourceKind, ProjectID: "test-project"},
			Client:  &http.Client{},
			BaseURL: cloudgdasrc.Endpoint,
		},
	}

	tcs := []struct {
		desc string
		cfg  cloudgdatool.Config
	}{
		{
			desc: "successful initialization",
			cfg: cloudgdatool.Config{
				Name:        "my-gda-query-tool",
				Kind:        "cloud-gemini-data-analytics-query",
				Source:      "gda-api-source",
				Description: "Test Description",
				Location:    "us-central1",
			},
		},
	}

	// Add an incompatible source for testing
	srcs["incompatible-source"] = &mockSource{kind: "another-kind"}

	for _, tc := range tcs {
		tc := tc
		t.Run(tc.desc, func(t *testing.T) {
			t.Parallel()
			tool, err := tc.cfg.Initialize(srcs)
			if err != nil {
				t.Fatalf("did not expect an error but got: %v", err)
			}
			// Basic sanity check on the returned tool
			_ = tool // Avoid unused variable error
		})
	}
}

func TestInvoke(t *testing.T) {
	t.Parallel()
	// Mock the HTTP client and server for Invoke testing
	serverMux := http.NewServeMux()
	// Update expected URL path to include the location "us-central1"
	serverMux.HandleFunc("/v1beta/projects/test-project/locations/global:queryData", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			t.Errorf("expected POST method, got %s", r.Method)
			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
		if r.Header.Get("Content-Type") != "application/json" {
			t.Errorf("expected Content-Type application/json, got %s", r.Header.Get("Content-Type"))
			http.Error(w, "Bad request", http.StatusBadRequest)
			return
		}

		// Read and unmarshal the request body
		bodyBytes, err := io.ReadAll(r.Body)
		if err != nil {
			t.Errorf("failed to read request body: %v", err)
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
			return
		}
		var reqPayload cloudgdatool.QueryDataRequest
		if err := json.Unmarshal(bodyBytes, &reqPayload); err != nil {
			t.Errorf("failed to unmarshal request payload: %v", err)
			http.Error(w, "Bad request", http.StatusBadRequest)
			return
		}

		// Verify expected fields
		if r.Header.Get("Authorization") == "" {
			t.Errorf("expected Authorization header, got empty")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}
		if reqPayload.Prompt != "How many accounts who have region in Prague are eligible for loans?" {
			t.Errorf("unexpected prompt: %s", reqPayload.Prompt)
		}

		// Verify payload's parent uses the tool's configured location
		if reqPayload.Parent != fmt.Sprintf("projects/%s/locations/%s", "test-project", "us-central1") {
			t.Errorf("unexpected payload parent: got %q, want %q", reqPayload.Parent, fmt.Sprintf("projects/%s/locations/%s", "test-project", "us-central1"))
		}

		// Verify context from config
		if reqPayload.Context == nil ||
			reqPayload.Context.DatasourceReferences == nil ||
			reqPayload.Context.DatasourceReferences.SpannerReference == nil ||
			reqPayload.Context.DatasourceReferences.SpannerReference.DatabaseReference == nil ||
			reqPayload.Context.DatasourceReferences.SpannerReference.DatabaseReference.ProjectID != "cloud-db-nl2sql" {
			t.Errorf("unexpected context: %v", reqPayload.Context)
		}

		// Verify generation options from config
		if reqPayload.GenerationOptions == nil || !reqPayload.GenerationOptions.GenerateQueryResult {
			t.Errorf("unexpected generation options: %v", reqPayload.GenerationOptions)
		}

		// Simulate a successful response
		resp := map[string]any{
			"queryResult":           "SELECT count(*) FROM accounts WHERE region = 'Prague' AND eligible_for_loans = true;",
			"naturalLanguageAnswer": "There are 5 accounts in Prague eligible for loans.",
		}
		_ = json.NewEncoder(w).Encode(resp)
	})

	mockServer := httptest.NewServer(serverMux)
	defer mockServer.Close()

	ctx := testutils.ContextWithUserAgent(context.Background(), "test-user-agent")

	// Create an authenticated client that uses the mock server
	authTransport := &authRoundTripper{
		Token: "Bearer test-access-token",
		Next:  mockServer.Client().Transport,
	}
	authClient := &http.Client{Transport: authTransport}

	// Create a real cloudgdasrc.Source but inject the authenticated client
	mockGdaSource := &cloudgdasrc.Source{
		Config:  cloudgdasrc.Config{Name: "mock-gda-source", Kind: cloudgdasrc.SourceKind, ProjectID: "test-project"},
		Client:  authClient,
		BaseURL: mockServer.URL,
	}
	srcs := map[string]sources.Source{
		"mock-gda-source": mockGdaSource,
	}

	// Initialize the tool config with context
	toolCfg := cloudgdatool.Config{
		Name:        "query-data-tool",
		Kind:        "cloud-gemini-data-analytics-query",
		Source:      "mock-gda-source",
		Description: "Query Gemini Data Analytics",
		Location:    "us-central1", // Set location for the test
		Context: &cloudgdatool.QueryDataContext{
			DatasourceReferences: &cloudgdatool.DatasourceReferences{
				SpannerReference: &cloudgdatool.SpannerReference{
					DatabaseReference: &cloudgdatool.SpannerDatabaseReference{
						ProjectID:  "cloud-db-nl2sql",
						Region:     "us-central1",
						InstanceID: "evalbench",
						DatabaseID: "financial",
						Engine:     cloudgdatool.SpannerEngineGoogleSQL,
					},
					AgentContextReference: &cloudgdatool.AgentContextReference{
						ContextSetID: "projects/cloud-db-nl2sql/locations/us-east1/contextSets/bdf_gsql_gemini_all_templates",
					},
				},
			},
		},
		GenerationOptions: &cloudgdatool.GenerationOptions{
			GenerateQueryResult: true,
		},
	}

	tool, err := toolCfg.Initialize(srcs)
	if err != nil {
		t.Fatalf("failed to initialize tool: %v", err)
	}

	// Prepare parameters for invocation - ONLY prompt
	params := parameters.ParamValues{
		{Name: "prompt", Value: "How many accounts who have region in Prague are eligible for loans?"},
	}

	resourceMgr := resources.NewResourceManager(srcs, nil, nil, nil, nil, nil, nil)

	// Invoke the tool
	result, err := tool.Invoke(ctx, resourceMgr, params, "") // No accessToken needed for ADC client
	if err != nil {
		t.Fatalf("tool invocation failed: %v", err)
	}

	// Validate the result
	expectedResult := map[string]any{
		"queryResult":           "SELECT count(*) FROM accounts WHERE region = 'Prague' AND eligible_for_loans = true;",
		"naturalLanguageAnswer": "There are 5 accounts in Prague eligible for loans.",
	}

	if !cmp.Equal(expectedResult, result) {
		t.Errorf("unexpected result: got %v, want %v", result, expectedResult)
	}
}