chore: update dataplex aspecttypes integration tests (#2193)

## Description

Addresses an issue where Dataplex AspectTypes created during integration
tests were not consistently deleted. This accumulation led to exceeding
the AspectType quota.
To prevent this, the test setup now includes a step to list and delete
all existing AspectTypes within the test project and location *before*
attempting to create any new ones. This ensures a clean state for each
test run and avoids hitting the quota.


## PR Checklist

> Thank you for opening a Pull Request! Before submitting your PR, there
are a
> few things you can do to make sure it goes smoothly:

- [ ] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [ ] Make sure to open an issue as a

[bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose)
  before writing your code! That way we can discuss the change, evaluate
  designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #2057

Co-authored-by: Averi Kitsch <akitsch@google.com>

cmd/root.go

@@ -313,15 +313,15 @@ func Execute() {
 type Command struct {
 	*cobra.Command
 
-	cfg             server.ServerConfig
-	logger          log.Logger
-	tools_file      string
-	tools_files     []string
-	tools_folder    string
-	prebuiltConfigs []string
-	inStream        io.Reader
-	outStream       io.Writer
-	errStream       io.Writer
+	cfg            server.ServerConfig
+	logger         log.Logger
+	tools_file     string
+	tools_files    []string
+	tools_folder   string
+	prebuiltConfig string
+	inStream       io.Reader
+	outStream      io.Writer
+	errStream      io.Writer
 }
 
 // NewCommand returns a Command object representing an invocation of the CLI.
@@ -374,10 +374,10 @@ func NewCommand(opts ...Option) *Command {
 	flags.StringVar(&cmd.cfg.TelemetryServiceName, "telemetry-service-name", "toolbox", "Sets the value of the service.name resource attribute for telemetry data.")
 	// Fetch prebuilt tools sources to customize the help description
 	prebuiltHelp := fmt.Sprintf(
-		"Use a prebuilt tool configuration by source type. Allowed: '%s'. Can be specified multiple times.",
+		"Use a prebuilt tool configuration by source type. Allowed: '%s'.",
 		strings.Join(prebuiltconfigs.GetPrebuiltSources(), "', '"),
 	)
-	flags.StringSliceVar(&cmd.prebuiltConfigs, "prebuilt", []string{}, prebuiltHelp)
+	flags.StringVar(&cmd.prebuiltConfig, "prebuilt", "", prebuiltHelp)
 	flags.BoolVar(&cmd.cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.")
 	flags.BoolVar(&cmd.cfg.DisableReload, "disable-reload", false, "Disables dynamic reloading of tools file.")
 	flags.BoolVar(&cmd.cfg.UI, "ui", false, "Launches the Toolbox UI web server.")
@@ -864,32 +864,24 @@ func run(cmd *Command) error {
 	var allToolsFiles []ToolsFile
 
 	// Load Prebuilt Configuration
-
-	if len(cmd.prebuiltConfigs) > 0 {
-		slices.Sort(cmd.prebuiltConfigs)
-		sourcesList := strings.Join(cmd.prebuiltConfigs, ", ")
-		logMsg := fmt.Sprintf("Using prebuilt tool configurations for: %s", sourcesList)
-		cmd.logger.InfoContext(ctx, logMsg)
-
-		for _, configName := range cmd.prebuiltConfigs {
-			buf, err := prebuiltconfigs.Get(configName)
-			if err != nil {
-				cmd.logger.ErrorContext(ctx, err.Error())
-				return err
-			}
-
-			// Update version string
-			cmd.cfg.Version += "+prebuilt." + configName
-
-			// Parse into ToolsFile struct
-			parsed, err := parseToolsFile(ctx, buf)
-			if err != nil {
-				errMsg := fmt.Errorf("unable to parse prebuilt tool configuration for '%s': %w", configName, err)
-				cmd.logger.ErrorContext(ctx, errMsg.Error())
-				return errMsg
-			}
-			allToolsFiles = append(allToolsFiles, parsed)
+	if cmd.prebuiltConfig != "" {
+		buf, err := prebuiltconfigs.Get(cmd.prebuiltConfig)
+		if err != nil {
+			cmd.logger.ErrorContext(ctx, err.Error())
+			return err
 		}
+		logMsg := fmt.Sprint("Using prebuilt tool configuration for ", cmd.prebuiltConfig)
+		cmd.logger.InfoContext(ctx, logMsg)
+		// Append prebuilt.source to Version string for the User Agent
+		cmd.cfg.Version += "+prebuilt." + cmd.prebuiltConfig
+
+		parsed, err := parseToolsFile(ctx, buf)
+		if err != nil {
+			errMsg := fmt.Errorf("unable to parse prebuilt tool configuration: %w", err)
+			cmd.logger.ErrorContext(ctx, errMsg.Error())
+			return errMsg
+		}
+		allToolsFiles = append(allToolsFiles, parsed)
 	}
 
 	// Determine if Custom Files should be loaded
@@ -897,7 +889,7 @@ func run(cmd *Command) error {
 	isCustomConfigured := cmd.tools_file != "" || len(cmd.tools_files) > 0 || cmd.tools_folder != ""
 
 	// Determine if default 'tools.yaml' should be used (No prebuilt AND No custom flags)
-	useDefaultToolsFile := len(cmd.prebuiltConfigs) == 0 && !isCustomConfigured
+	useDefaultToolsFile := cmd.prebuiltConfig == "" && !isCustomConfigured
 
 	if useDefaultToolsFile {
 		cmd.tools_file = "tools.yaml"

cmd/root_test.go

@@ -420,27 +420,17 @@ func TestPrebuiltFlag(t *testing.T) {
 	tcs := []struct {
 		desc string
 		args []string
-		want []string
+		want string
 	}{
 		{
 			desc: "default value",
 			args: []string{},
-			want: []string{},
+			want: "",
 		},
 		{
-			desc: "single prebuilt flag",
-			args: []string{"--prebuilt", "alloydb"},
-			want: []string{"alloydb"},
-		},
-		{
-			desc: "multiple prebuilt flags",
-			args: []string{"--prebuilt", "alloydb", "--prebuilt", "bigquery"},
-			want: []string{"alloydb", "bigquery"},
-		},
-		{
-			desc: "comma separated prebuilt flags",
-			args: []string{"--prebuilt", "alloydb,bigquery"},
-			want: []string{"alloydb", "bigquery"},
+			desc: "custom pre built flag",
+			args: []string{"--tools-file", "alloydb"},
+			want: "alloydb",
 		},
 	}
 	for _, tc := range tcs {
@@ -449,8 +439,8 @@ func TestPrebuiltFlag(t *testing.T) {
 			if err != nil {
 				t.Fatalf("unexpected error invoking command: %s", err)
 			}
-			if diff := cmp.Diff(c.prebuiltConfigs, tc.want); diff != "" {
-				t.Fatalf("got %v, want %v, diff %s", c.prebuiltConfigs, tc.want, diff)
+			if c.tools_file != tc.want {
+				t.Fatalf("got %v, want %v", c.cfg, tc.want)
 			}
 		})
 	}
@@ -2073,12 +2063,6 @@ authSources:
 				return nil
 			},
 		},
-		{
-			desc:      "sqlite called twice error",
-			args:      []string{"--prebuilt", "sqlite", "--prebuilt", "sqlite"},
-			wantErr:   true,
-			errString: "resource conflicts detected",
-		},
 		{
 			desc:      "tool conflict error",
 			args:      []string{"--prebuilt", "sqlite", "--tools-file", toolConflictFile},

docs/en/reference/cli.md

@@ -16,7 +16,7 @@ description: >
 |              | `--log-level`              | Specify the minimum level logged. Allowed: 'DEBUG', 'INFO', 'WARN', 'ERROR'.                                                                                                     | `info`      |
 |              | `--logging-format`         | Specify logging format to use. Allowed: 'standard' or 'JSON'.                                                                                                                    | `standard`  |
 | `-p`         | `--port`                   | Port the server will listen on.                                                                                                                                                  | `5000`      |
-|              | `--prebuilt`               | Use one or more prebuilt tool configuration by source type. See [Prebuilt Tools Reference](prebuilt-tools.md) for allowed values.                                                          |             |
+|              | `--prebuilt`               | Use a prebuilt tool configuration by source type. See [Prebuilt Tools Reference](prebuilt-tools.md) for allowed values.                                                          |             |
 |              | `--stdio`                  | Listens via MCP STDIO instead of acting as a remote HTTP server.                                                                                                                 |             |
 |              | `--telemetry-gcp`          | Enable exporting directly to Google Cloud Monitoring.                                                                                                                            |             |
 |              | `--telemetry-otlp`         | Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318')                                                                    |             |
@@ -50,11 +50,6 @@ description: >
 
 # Server with prebuilt + custom tools configurations
 ./toolbox --tools-file tools.yaml --prebuilt alloydb-postgres
-
-# Server with multiple prebuilt tools configurations
-./toolbox --prebuilt alloydb-postgres,alloydb-postgres-admin
-# OR
-./toolbox --prebuilt alloydb-postgres --prebuilt alloydb-postgres-admin
 ```
 
 ### Tool Configuration Sources
@@ -75,7 +70,7 @@ The CLI supports multiple mutually exclusive ways to specify tool configurations
 
 **Prebuilt Configurations:**
 
-- `--prebuilt`: Use one or more predefined configurations for specific database types (e.g.,
+- `--prebuilt`: Use predefined configurations for specific database types (e.g.,
   'bigquery', 'postgres', 'spanner'). See [Prebuilt Tools
   Reference](prebuilt-tools.md) for allowed values.
 

docs/en/reference/prebuilt-tools.md

@@ -16,9 +16,6 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
 {{< notice tip >}}
 You can now use `--prebuilt` along `--tools-file`, `--tools-files`, or
 `--tools-folder` to combine prebuilt configs with custom tools.
-
-You can also combine multiple prebuilt configs.
-
 See [Usage Examples](../reference/cli.md#examples).
 {{< /notice >}}
 

tests/dataplex/dataplex_integration_test.go

@@ -85,13 +85,66 @@ func initDataplexConnection(ctx context.Context) (*dataplex.CatalogClient, error
 	return client, nil
 }
 
+// cleanupOldAspectTypes Deletes AspectTypes older than the specified duration.
+func cleanupOldAspectTypes(t *testing.T, ctx context.Context, client *dataplex.CatalogClient, oldThreshold time.Duration) {
+	parent := fmt.Sprintf("projects/%s/locations/us", DataplexProject)
+	olderThanTime := time.Now().Add(-oldThreshold)
+
+	listReq := &dataplexpb.ListAspectTypesRequest{
+		Parent:   parent,
+		PageSize: 100,               // Fetch up to 100 items
+		OrderBy:  "create_time asc", // Order by creation time
+	}
+
+	const maxDeletes = 8 // Explicitly limit the number of deletions
+	it := client.ListAspectTypes(ctx, listReq)
+	var aspectTypesToDelete []string
+	for len(aspectTypesToDelete) < maxDeletes {
+		aspectType, err := it.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			t.Logf("Warning: Failed to list aspect types during cleanup: %v", err)
+			return
+		}
+		// Perform time-based filtering in memory
+		if aspectType.CreateTime != nil {
+			createTime := aspectType.CreateTime.AsTime()
+			if createTime.Before(olderThanTime) {
+				aspectTypesToDelete = append(aspectTypesToDelete, aspectType.GetName())
+			}
+		} else {
+			t.Logf("Warning: AspectType %s has no CreateTime", aspectType.GetName())
+		}
+	}
+	if len(aspectTypesToDelete) == 0 {
+		t.Logf("cleanupOldAspectTypes: No aspect types found older than %s to delete.", oldThreshold.String())
+		return
+	}
+
+	for _, aspectTypeName := range aspectTypesToDelete {
+		deleteReq := &dataplexpb.DeleteAspectTypeRequest{Name: aspectTypeName}
+		op, err := client.DeleteAspectType(ctx, deleteReq)
+		if err != nil {
+			t.Logf("Warning: Failed to delete aspect type %s: %v", aspectTypeName, err)
+			continue // Skip to the next item if initiation fails
+		}
+
+		if err := op.Wait(ctx); err != nil {
+			t.Logf("Warning: Failed to delete aspect type %s, operation error: %v", aspectTypeName, err)
+		} else {
+			t.Logf("cleanupOldAspectTypes: Successfully deleted %s", aspectTypeName)
+		}
+	}
+}
+
 func TestDataplexToolEndpoints(t *testing.T) {
 	sourceConfig := getDataplexVars(t)
 	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
 	defer cancel()
 
 	var args []string
-
 	bigqueryClient, err := initBigQueryConnection(ctx, DataplexProject)
 	if err != nil {
 		t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
@@ -102,6 +155,9 @@ func TestDataplexToolEndpoints(t *testing.T) {
 		t.Fatalf("unable to create Dataplex connection: %s", err)
 	}
 
+	// Cleanup older aspecttypes
+	cleanupOldAspectTypes(t, ctx, dataplexClient, 1*time.Hour)
+
 	// create resources with UUID
 	datasetName := fmt.Sprintf("temp_toolbox_test_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
 	tableName := fmt.Sprintf("param_table_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))