googleapis--mcp-toolbox/cmd/internal/migrate/command_test.go

// Copyright 2026 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 migrate

import (
	"bytes"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/googleapis/mcp-toolbox/cmd/internal"
	"github.com/spf13/cobra"
)

func invokeCommand(args []string) (string, error) {
	parentCmd := &cobra.Command{Use: "toolbox"}

	buf := new(bytes.Buffer)
	opts := internal.NewToolboxOptions(internal.WithIOStreams(buf, buf))
	internal.PersistentFlags(parentCmd, opts)

	cmd := NewCommand(opts)
	parentCmd.AddCommand(cmd)
	parentCmd.SetArgs(args)

	err := parentCmd.Execute()
	return buf.String(), err
}

func TestMigrate(t *testing.T) {
	toolsFileContent := `
sources:
  my-pg-instance:
    kind: cloud-sql-postgres
    project: my-project
    region: my-region
    instance: my-instance
    database: my_db
    user: my_user
    password: my_pass`

	toolsFileContentNew := `
kind: source
name: my-pg-instance
type: cloud-sql-postgres
project: my-project
region: my-region
instance: my-instance
database: my_db
user: my_user
password: my_pass
`
	toolsFileContent2 := `
tools:
  example_tool2:
    kind: postgres-sql
    source: my-pg-instance
    description: some description
    statement: SELECT * FROM SQL_STATEMENT;
    parameters:
      - name: country
        type: string
        description: some description`

	toolsFileContent2New := `
kind: tool
name: example_tool2
type: postgres-sql
source: my-pg-instance
description: some description
statement: SELECT * FROM SQL_STATEMENT;
parameters:
- name: country
  type: string
  description: some description
`

	t.Run("migrate tools file", func(t *testing.T) {
		tmpDir := t.TempDir()

		toolsFilePath := filepath.Join(tmpDir, "foo.yaml")
		if err := os.WriteFile(toolsFilePath, []byte(toolsFileContent), 0644); err != nil {
			t.Fatalf("failed to write tools file: %v", err)
		}

		args := []string{"migrate", "--tools-file", toolsFilePath}
		got, err := invokeCommand(args)
		if err != nil {
			t.Fatalf("command failed: %v\nOutput: %s", err, got)
		}

		// verify backup file
		backupFile := toolsFilePath + ".bak"
		_, err = os.Stat(backupFile)
		if err != nil {
			t.Fatalf("error verifying backup file: %v", err)
		}
		actualContent, err := os.ReadFile(backupFile)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent, actualContent)
		}

		// check content of new file
		actualContent, err = os.ReadFile(toolsFilePath)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContentNew)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContentNew, actualContent)
		}
	})
	t.Run("migrate tools files", func(t *testing.T) {
		tmpDir := t.TempDir()

		toolsFilePath1 := filepath.Join(tmpDir, "foo.yaml")
		if err := os.WriteFile(toolsFilePath1, []byte(toolsFileContent), 0644); err != nil {
			t.Fatalf("failed to write tools file: %v", err)
		}
		toolsFilePath2 := filepath.Join(tmpDir, "foo2.yaml")
		if err := os.WriteFile(toolsFilePath2, []byte(toolsFileContent2), 0644); err != nil {
			t.Fatalf("failed to write tools file: %v", err)
		}

		toolsFiles := toolsFilePath1 + "," + toolsFilePath2
		args := []string{"migrate", "--tools-files", toolsFiles}
		got, err := invokeCommand(args)
		if err != nil {
			t.Fatalf("command failed: %v\nOutput: %s", err, got)
		}

		// verify backup file1
		backupFile := toolsFilePath1 + ".bak"
		_, err = os.Stat(backupFile)
		if err != nil {
			t.Fatalf("error verifying backup file: %v", err)
		}
		actualContent, err := os.ReadFile(backupFile)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent, actualContent)
		}

		// verify backup file2
		backupFile = toolsFilePath2 + ".bak"
		_, err = os.Stat(backupFile)
		if err != nil {
			t.Fatalf("error verifying backup file: %v", err)
		}
		actualContent, err = os.ReadFile(backupFile)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent2)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent2, actualContent)
		}

		// check content of new file1
		actualContent, err = os.ReadFile(toolsFilePath1)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContentNew)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContentNew, actualContent)
		}

		// check content of new file2
		actualContent, err = os.ReadFile(toolsFilePath2)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent2New)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent2New, actualContent)
		}
	})

	t.Run("migrate tools folder", func(t *testing.T) {
		tmpDir := t.TempDir()

		toolsFilePath1 := filepath.Join(tmpDir, "foo.yaml")
		if err := os.WriteFile(toolsFilePath1, []byte(toolsFileContent), 0644); err != nil {
			t.Fatalf("failed to write tools file: %v", err)
		}
		toolsFilePath2 := filepath.Join(tmpDir, "foo2.yaml")
		if err := os.WriteFile(toolsFilePath2, []byte(toolsFileContent2), 0644); err != nil {
			t.Fatalf("failed to write tools file: %v", err)
		}

		args := []string{"migrate", "--tools-folder", tmpDir}
		got, err := invokeCommand(args)
		if err != nil {
			t.Fatalf("command failed: %v\nOutput: %s", err, got)
		}

		// verify backup file1
		backupFile := toolsFilePath1 + ".bak"
		_, err = os.Stat(backupFile)
		if err != nil {
			t.Fatalf("error verifying backup file: %v", err)
		}
		actualContent, err := os.ReadFile(backupFile)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent, actualContent)
		}

		// verify backup file2
		backupFile = toolsFilePath2 + ".bak"
		_, err = os.Stat(backupFile)
		if err != nil {
			t.Fatalf("error verifying backup file: %v", err)
		}
		actualContent, err = os.ReadFile(backupFile)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent2)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent2, actualContent)
		}

		// check content of new file1
		actualContent, err = os.ReadFile(toolsFilePath1)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContentNew)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContentNew, actualContent)
		}

		// check content of new file2
		actualContent, err = os.ReadFile(toolsFilePath2)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent2New)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent2New, actualContent)
		}
	})
}

func TestMigrateDryRun(t *testing.T) {
	toolsFileContent := `
sources:
  my-pg-instance:
    kind: cloud-sql-postgres
    project: my-project
    region: my-region
    instance: my-instance
    database: my_db
    user: my_user
    password: my_pass
`

	toolsFileContentNew := `
kind: source
name: my-pg-instance
type: cloud-sql-postgres
project: my-project
region: my-region
instance: my-instance
database: my_db
user: my_user
password: my_pass
`
	toolsFileContent2 := `
tools:
  example_tool2:
    kind: postgres-sql
    source: my-pg-instance
    description: |
      some description
      - string
    statement: SELECT * FROM SQL_STATEMENT;
    parameters:
      - name: country
        type: string
        description: some description`

	toolsFileContent2New := `
kind: tool
name: example_tool2
type: postgres-sql
source: my-pg-instance
description: |
  some description
  - string
statement: SELECT * FROM SQL_STATEMENT;
parameters:
- name: country
  type: string
  description: some description
`

	t.Run("migrate tools file", func(t *testing.T) {
		tmpDir := t.TempDir()

		toolsFilePath := filepath.Join(tmpDir, "foo.yaml")
		if err := os.WriteFile(toolsFilePath, []byte(toolsFileContent), 0644); err != nil {
			t.Fatalf("failed to write tools file: %v", err)
		}

		args := []string{"migrate", "--tools-file", toolsFilePath, "--dry-run"}
		got, err := invokeCommand(args)
		if err != nil {
			t.Fatalf("command failed: %v\nOutput: %s", err, got)
		}

		// verify original file
		actualContent, err := os.ReadFile(toolsFilePath)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent, actualContent)
		}

		// check output
		if !strings.Contains(got, toolsFileContentNew) {
			t.Fatalf("expected output not found!\nExpected: %q\nGot: %q", toolsFileContentNew, got)
		}
	})
	t.Run("migrate tools files", func(t *testing.T) {
		tmpDir := t.TempDir()

		toolsFilePath1 := filepath.Join(tmpDir, "foo.yaml")
		if err := os.WriteFile(toolsFilePath1, []byte(toolsFileContent), 0644); err != nil {
			t.Fatalf("failed to write tools file: %v", err)
		}
		toolsFilePath2 := filepath.Join(tmpDir, "foo2.yaml")
		if err := os.WriteFile(toolsFilePath2, []byte(toolsFileContent2), 0644); err != nil {
			t.Fatalf("failed to write tools file: %v", err)
		}

		toolsFiles := toolsFilePath1 + "," + toolsFilePath2
		args := []string{"migrate", "--tools-files", toolsFiles, "--dry-run"}
		got, err := invokeCommand(args)
		if err != nil {
			t.Fatalf("command failed: %v\nOutput: %s", err, got)
		}

		// verify original file1
		actualContent, err := os.ReadFile(toolsFilePath1)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent, actualContent)
		}

		// verify original file2
		actualContent, err = os.ReadFile(toolsFilePath2)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent2)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent2, actualContent)
		}

		// check output
		if !strings.Contains(got, toolsFileContentNew) {
			t.Fatalf("expected output not found!\nExpected: %q\nGot: %q", toolsFileContentNew, got)
		}
		if !strings.Contains(got, toolsFileContent2New) {
			t.Fatalf("expected output not found!\nExpected: %q\nGot: %q", toolsFileContent2New, got)
		}
	})

	t.Run("migrate tools folder", func(t *testing.T) {
		tmpDir := t.TempDir()

		toolsFilePath1 := filepath.Join(tmpDir, "foo.yaml")
		if err := os.WriteFile(toolsFilePath1, []byte(toolsFileContent), 0644); err != nil {
			t.Fatalf("failed to write tools file: %v", err)
		}
		toolsFilePath2 := filepath.Join(tmpDir, "foo2.yaml")
		if err := os.WriteFile(toolsFilePath2, []byte(toolsFileContent2), 0644); err != nil {
			t.Fatalf("failed to write tools file: %v", err)
		}

		args := []string{"migrate", "--tools-folder", tmpDir, "--dry-run"}
		got, err := invokeCommand(args)
		if err != nil {
			t.Fatalf("command failed: %v\nOutput: %s", err, got)
		}

		// verify original file1
		actualContent, err := os.ReadFile(toolsFilePath1)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent, actualContent)
		}

		// verify original file2
		actualContent, err = os.ReadFile(toolsFilePath2)
		if err != nil {
			t.Fatalf("failed to read backup file: %v", err)
		}
		if !bytes.Equal(actualContent, []byte(toolsFileContent2)) {
			t.Fatalf("file content mismatch!\nExpected: %q\nGot: %q", toolsFileContent2, actualContent)
		}

		// check output
		if !strings.Contains(got, toolsFileContentNew) {
			t.Fatalf("expected output not found!\nExpected: %q\nGot: %q", toolsFileContentNew, got)
		}
		if !strings.Contains(got, toolsFileContent2New) {
			t.Fatalf("expected output not found!\nExpected: %q\nGot: %q", toolsFileContent2New, got)
		}
	})
}