This is page 31 of 33. Use http://codebase.md/googleapis/genai-toolbox?page={x} to view the full context. # Directory Structure ``` ├── .ci │ ├── continuous.release.cloudbuild.yaml │ ├── generate_release_table.sh │ ├── integration.cloudbuild.yaml │ ├── quickstart_test │ │ ├── go.integration.cloudbuild.yaml │ │ ├── js.integration.cloudbuild.yaml │ │ ├── py.integration.cloudbuild.yaml │ │ ├── run_go_tests.sh │ │ ├── run_js_tests.sh │ │ ├── run_py_tests.sh │ │ └── setup_hotels_sample.sql │ ├── test_with_coverage.sh │ └── versioned.release.cloudbuild.yaml ├── .github │ ├── auto-label.yaml │ ├── blunderbuss.yml │ ├── CODEOWNERS │ ├── header-checker-lint.yml │ ├── ISSUE_TEMPLATE │ │ ├── bug_report.yml │ │ ├── config.yml │ │ ├── feature_request.yml │ │ └── question.yml │ ├── label-sync.yml │ ├── labels.yaml │ ├── PULL_REQUEST_TEMPLATE.md │ ├── release-please.yml │ ├── renovate.json5 │ ├── sync-repo-settings.yaml │ └── workflows │ ├── cloud_build_failure_reporter.yml │ ├── deploy_dev_docs.yaml │ ├── deploy_previous_version_docs.yaml │ ├── deploy_versioned_docs.yaml │ ├── docs_deploy.yaml │ ├── docs_preview_clean.yaml │ ├── docs_preview_deploy.yaml │ ├── lint.yaml │ ├── schedule_reporter.yml │ ├── sync-labels.yaml │ └── tests.yaml ├── .gitignore ├── .gitmodules ├── .golangci.yaml ├── .hugo │ ├── archetypes │ │ └── default.md │ ├── assets │ │ ├── icons │ │ │ └── logo.svg │ │ └── scss │ │ ├── _styles_project.scss │ │ └── _variables_project.scss │ ├── go.mod │ ├── go.sum │ ├── hugo.toml │ ├── layouts │ │ ├── _default │ │ │ └── home.releases.releases │ │ ├── index.llms-full.txt │ │ ├── index.llms.txt │ │ ├── partials │ │ │ ├── hooks │ │ │ │ └── head-end.html │ │ │ ├── navbar-version-selector.html │ │ │ ├── page-meta-links.html │ │ │ └── td │ │ │ └── render-heading.html │ │ ├── robot.txt │ │ └── shortcodes │ │ ├── include.html │ │ ├── ipynb.html │ │ └── regionInclude.html │ ├── package-lock.json │ ├── package.json │ └── static │ ├── favicons │ │ ├── android-chrome-192x192.png │ │ ├── android-chrome-512x512.png │ │ ├── apple-touch-icon.png │ │ ├── favicon-16x16.png │ │ ├── favicon-32x32.png │ │ └── favicon.ico │ └── js │ └── w3.js ├── CHANGELOG.md ├── cmd │ ├── options_test.go │ ├── options.go │ ├── root_test.go │ ├── root.go │ └── version.txt ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── DEVELOPER.md ├── Dockerfile ├── docs │ └── en │ ├── _index.md │ ├── about │ │ ├── _index.md │ │ └── faq.md │ ├── concepts │ │ ├── _index.md │ │ └── telemetry │ │ ├── index.md │ │ ├── telemetry_flow.png │ │ └── telemetry_traces.png │ ├── getting-started │ │ ├── _index.md │ │ ├── colab_quickstart.ipynb │ │ ├── configure.md │ │ ├── introduction │ │ │ ├── _index.md │ │ │ └── architecture.png │ │ ├── local_quickstart_go.md │ │ ├── local_quickstart_js.md │ │ ├── local_quickstart.md │ │ ├── mcp_quickstart │ │ │ ├── _index.md │ │ │ ├── inspector_tools.png │ │ │ └── inspector.png │ │ └── quickstart │ │ ├── go │ │ │ ├── genAI │ │ │ │ ├── go.mod │ │ │ │ ├── go.sum │ │ │ │ └── quickstart.go │ │ │ ├── genkit │ │ │ │ ├── go.mod │ │ │ │ ├── go.sum │ │ │ │ └── quickstart.go │ │ │ ├── langchain │ │ │ │ ├── go.mod │ │ │ │ ├── go.sum │ │ │ │ └── quickstart.go │ │ │ ├── openAI │ │ │ │ ├── go.mod │ │ │ │ ├── go.sum │ │ │ │ └── quickstart.go │ │ │ └── quickstart_test.go │ │ ├── golden.txt │ │ ├── js │ │ │ ├── genAI │ │ │ │ ├── package-lock.json │ │ │ │ ├── package.json │ │ │ │ └── quickstart.js │ │ │ ├── genkit │ │ │ │ ├── package-lock.json │ │ │ │ ├── package.json │ │ │ │ └── quickstart.js │ │ │ ├── langchain │ │ │ │ ├── package-lock.json │ │ │ │ ├── package.json │ │ │ │ └── quickstart.js │ │ │ ├── llamaindex │ │ │ │ ├── package-lock.json │ │ │ │ ├── package.json │ │ │ │ └── quickstart.js │ │ │ └── quickstart.test.js │ │ ├── python │ │ │ ├── __init__.py │ │ │ ├── adk │ │ │ │ ├── quickstart.py │ │ │ │ └── requirements.txt │ │ │ ├── core │ │ │ │ ├── quickstart.py │ │ │ │ └── requirements.txt │ │ │ ├── langchain │ │ │ │ ├── quickstart.py │ │ │ │ └── requirements.txt │ │ │ ├── llamaindex │ │ │ │ ├── quickstart.py │ │ │ │ └── requirements.txt │ │ │ └── quickstart_test.py │ │ └── shared │ │ ├── cloud_setup.md │ │ ├── configure_toolbox.md │ │ └── database_setup.md │ ├── how-to │ │ ├── _index.md │ │ ├── connect_via_geminicli.md │ │ ├── connect_via_mcp.md │ │ ├── connect-ide │ │ │ ├── _index.md │ │ │ ├── alloydb_pg_admin_mcp.md │ │ │ ├── alloydb_pg_mcp.md │ │ │ ├── bigquery_mcp.md │ │ │ ├── cloud_sql_mssql_admin_mcp.md │ │ │ ├── cloud_sql_mssql_mcp.md │ │ │ ├── cloud_sql_mysql_admin_mcp.md │ │ │ ├── cloud_sql_mysql_mcp.md │ │ │ ├── cloud_sql_pg_admin_mcp.md │ │ │ ├── cloud_sql_pg_mcp.md │ │ │ ├── firestore_mcp.md │ │ │ ├── looker_mcp.md │ │ │ ├── mssql_mcp.md │ │ │ ├── mysql_mcp.md │ │ │ ├── neo4j_mcp.md │ │ │ ├── postgres_mcp.md │ │ │ ├── spanner_mcp.md │ │ │ └── sqlite_mcp.md │ │ ├── deploy_docker.md │ │ ├── deploy_gke.md │ │ ├── deploy_toolbox.md │ │ ├── export_telemetry.md │ │ └── toolbox-ui │ │ ├── edit-headers.gif │ │ ├── edit-headers.png │ │ ├── index.md │ │ ├── optional-param-checked.png │ │ ├── optional-param-unchecked.png │ │ ├── run-tool.gif │ │ ├── tools.png │ │ └── toolsets.png │ ├── reference │ │ ├── _index.md │ │ ├── cli.md │ │ └── prebuilt-tools.md │ ├── resources │ │ ├── _index.md │ │ ├── authServices │ │ │ ├── _index.md │ │ │ └── google.md │ │ ├── sources │ │ │ ├── _index.md │ │ │ ├── alloydb-admin.md │ │ │ ├── alloydb-pg.md │ │ │ ├── bigquery.md │ │ │ ├── bigtable.md │ │ │ ├── cassandra.md │ │ │ ├── clickhouse.md │ │ │ ├── cloud-monitoring.md │ │ │ ├── cloud-sql-admin.md │ │ │ ├── cloud-sql-mssql.md │ │ │ ├── cloud-sql-mysql.md │ │ │ ├── cloud-sql-pg.md │ │ │ ├── couchbase.md │ │ │ ├── dataplex.md │ │ │ ├── dgraph.md │ │ │ ├── firebird.md │ │ │ ├── firestore.md │ │ │ ├── http.md │ │ │ ├── looker.md │ │ │ ├── mongodb.md │ │ │ ├── mssql.md │ │ │ ├── mysql.md │ │ │ ├── neo4j.md │ │ │ ├── oceanbase.md │ │ │ ├── oracle.md │ │ │ ├── postgres.md │ │ │ ├── redis.md │ │ │ ├── spanner.md │ │ │ ├── sqlite.md │ │ │ ├── tidb.md │ │ │ ├── trino.md │ │ │ ├── valkey.md │ │ │ └── yugabytedb.md │ │ └── tools │ │ ├── _index.md │ │ ├── alloydb │ │ │ ├── _index.md │ │ │ ├── alloydb-create-cluster.md │ │ │ ├── alloydb-create-instance.md │ │ │ ├── alloydb-create-user.md │ │ │ ├── alloydb-get-cluster.md │ │ │ ├── alloydb-get-instance.md │ │ │ ├── alloydb-get-user.md │ │ │ ├── alloydb-list-clusters.md │ │ │ ├── alloydb-list-instances.md │ │ │ ├── alloydb-list-users.md │ │ │ └── alloydb-wait-for-operation.md │ │ ├── alloydbainl │ │ │ ├── _index.md │ │ │ └── alloydb-ai-nl.md │ │ ├── bigquery │ │ │ ├── _index.md │ │ │ ├── bigquery-analyze-contribution.md │ │ │ ├── bigquery-conversational-analytics.md │ │ │ ├── bigquery-execute-sql.md │ │ │ ├── bigquery-forecast.md │ │ │ ├── bigquery-get-dataset-info.md │ │ │ ├── bigquery-get-table-info.md │ │ │ ├── bigquery-list-dataset-ids.md │ │ │ ├── bigquery-list-table-ids.md │ │ │ ├── bigquery-search-catalog.md │ │ │ └── bigquery-sql.md │ │ ├── bigtable │ │ │ ├── _index.md │ │ │ └── bigtable-sql.md │ │ ├── cassandra │ │ │ ├── _index.md │ │ │ └── cassandra-cql.md │ │ ├── clickhouse │ │ │ ├── _index.md │ │ │ ├── clickhouse-execute-sql.md │ │ │ ├── clickhouse-list-databases.md │ │ │ ├── clickhouse-list-tables.md │ │ │ └── clickhouse-sql.md │ │ ├── cloudmonitoring │ │ │ ├── _index.md │ │ │ └── cloud-monitoring-query-prometheus.md │ │ ├── cloudsql │ │ │ ├── _index.md │ │ │ ├── cloudsqlcreatedatabase.md │ │ │ ├── cloudsqlcreateusers.md │ │ │ ├── cloudsqlgetinstances.md │ │ │ ├── cloudsqllistdatabases.md │ │ │ ├── cloudsqllistinstances.md │ │ │ ├── cloudsqlmssqlcreateinstance.md │ │ │ ├── cloudsqlmysqlcreateinstance.md │ │ │ ├── cloudsqlpgcreateinstances.md │ │ │ └── cloudsqlwaitforoperation.md │ │ ├── couchbase │ │ │ ├── _index.md │ │ │ └── couchbase-sql.md │ │ ├── dataform │ │ │ ├── _index.md │ │ │ └── dataform-compile-local.md │ │ ├── dataplex │ │ │ ├── _index.md │ │ │ ├── dataplex-lookup-entry.md │ │ │ ├── dataplex-search-aspect-types.md │ │ │ └── dataplex-search-entries.md │ │ ├── dgraph │ │ │ ├── _index.md │ │ │ └── dgraph-dql.md │ │ ├── firebird │ │ │ ├── _index.md │ │ │ ├── firebird-execute-sql.md │ │ │ └── firebird-sql.md │ │ ├── firestore │ │ │ ├── _index.md │ │ │ ├── firestore-add-documents.md │ │ │ ├── firestore-delete-documents.md │ │ │ ├── firestore-get-documents.md │ │ │ ├── firestore-get-rules.md │ │ │ ├── firestore-list-collections.md │ │ │ ├── firestore-query-collection.md │ │ │ ├── firestore-query.md │ │ │ ├── firestore-update-document.md │ │ │ └── firestore-validate-rules.md │ │ ├── http │ │ │ ├── _index.md │ │ │ └── http.md │ │ ├── looker │ │ │ ├── _index.md │ │ │ ├── looker-add-dashboard-element.md │ │ │ ├── looker-conversational-analytics.md │ │ │ ├── looker-get-dashboards.md │ │ │ ├── looker-get-dimensions.md │ │ │ ├── looker-get-explores.md │ │ │ ├── looker-get-filters.md │ │ │ ├── looker-get-looks.md │ │ │ ├── looker-get-measures.md │ │ │ ├── looker-get-models.md │ │ │ ├── looker-get-parameters.md │ │ │ ├── looker-health-analyze.md │ │ │ ├── looker-health-pulse.md │ │ │ ├── looker-health-vacuum.md │ │ │ ├── looker-make-dashboard.md │ │ │ ├── looker-make-look.md │ │ │ ├── looker-query-sql.md │ │ │ ├── looker-query-url.md │ │ │ ├── looker-query.md │ │ │ └── looker-run-look.md │ │ ├── mongodb │ │ │ ├── _index.md │ │ │ ├── mongodb-aggregate.md │ │ │ ├── mongodb-delete-many.md │ │ │ ├── mongodb-delete-one.md │ │ │ ├── mongodb-find-one.md │ │ │ ├── mongodb-find.md │ │ │ ├── mongodb-insert-many.md │ │ │ ├── mongodb-insert-one.md │ │ │ ├── mongodb-update-many.md │ │ │ └── mongodb-update-one.md │ │ ├── mssql │ │ │ ├── _index.md │ │ │ ├── mssql-execute-sql.md │ │ │ ├── mssql-list-tables.md │ │ │ └── mssql-sql.md │ │ ├── mysql │ │ │ ├── _index.md │ │ │ ├── mysql-execute-sql.md │ │ │ ├── mysql-list-active-queries.md │ │ │ ├── mysql-list-table-fragmentation.md │ │ │ ├── mysql-list-tables-missing-unique-indexes.md │ │ │ ├── mysql-list-tables.md │ │ │ └── mysql-sql.md │ │ ├── neo4j │ │ │ ├── _index.md │ │ │ ├── neo4j-cypher.md │ │ │ ├── neo4j-execute-cypher.md │ │ │ └── neo4j-schema.md │ │ ├── oceanbase │ │ │ ├── _index.md │ │ │ ├── oceanbase-execute-sql.md │ │ │ └── oceanbase-sql.md │ │ ├── oracle │ │ │ ├── _index.md │ │ │ ├── oracle-execute-sql.md │ │ │ └── oracle-sql.md │ │ ├── postgres │ │ │ ├── _index.md │ │ │ ├── postgres-execute-sql.md │ │ │ ├── postgres-list-active-queries.md │ │ │ ├── postgres-list-available-extensions.md │ │ │ ├── postgres-list-installed-extensions.md │ │ │ ├── postgres-list-tables.md │ │ │ └── postgres-sql.md │ │ ├── redis │ │ │ ├── _index.md │ │ │ └── redis.md │ │ ├── spanner │ │ │ ├── _index.md │ │ │ ├── spanner-execute-sql.md │ │ │ ├── spanner-list-tables.md │ │ │ └── spanner-sql.md │ │ ├── sqlite │ │ │ ├── _index.md │ │ │ ├── sqlite-execute-sql.md │ │ │ └── sqlite-sql.md │ │ ├── tidb │ │ │ ├── _index.md │ │ │ ├── tidb-execute-sql.md │ │ │ └── tidb-sql.md │ │ ├── trino │ │ │ ├── _index.md │ │ │ ├── trino-execute-sql.md │ │ │ └── trino-sql.md │ │ ├── utility │ │ │ ├── _index.md │ │ │ └── wait.md │ │ ├── valkey │ │ │ ├── _index.md │ │ │ └── valkey.md │ │ └── yuagbytedb │ │ ├── _index.md │ │ └── yugabytedb-sql.md │ ├── samples │ │ ├── _index.md │ │ ├── alloydb │ │ │ ├── _index.md │ │ │ ├── ai-nl │ │ │ │ ├── alloydb_ai_nl.ipynb │ │ │ │ └── index.md │ │ │ └── mcp_quickstart.md │ │ ├── bigquery │ │ │ ├── _index.md │ │ │ ├── colab_quickstart_bigquery.ipynb │ │ │ ├── local_quickstart.md │ │ │ └── mcp_quickstart │ │ │ ├── _index.md │ │ │ ├── inspector_tools.png │ │ │ └── inspector.png │ │ └── looker │ │ ├── _index.md │ │ ├── looker_gemini_oauth │ │ │ ├── _index.md │ │ │ ├── authenticated.png │ │ │ ├── authorize.png │ │ │ └── registration.png │ │ ├── looker_gemini.md │ │ └── looker_mcp_inspector │ │ ├── _index.md │ │ ├── inspector_tools.png │ │ └── inspector.png │ └── sdks │ ├── _index.md │ ├── go-sdk.md │ ├── js-sdk.md │ └── python-sdk.md ├── go.mod ├── go.sum ├── internal │ ├── auth │ │ ├── auth.go │ │ └── google │ │ └── google.go │ ├── log │ │ ├── handler.go │ │ ├── log_test.go │ │ ├── log.go │ │ └── logger.go │ ├── prebuiltconfigs │ │ ├── prebuiltconfigs_test.go │ │ ├── prebuiltconfigs.go │ │ └── tools │ │ ├── alloydb-postgres-admin.yaml │ │ ├── alloydb-postgres-observability.yaml │ │ ├── alloydb-postgres.yaml │ │ ├── bigquery.yaml │ │ ├── clickhouse.yaml │ │ ├── cloud-sql-mssql-admin.yaml │ │ ├── cloud-sql-mssql-observability.yaml │ │ ├── cloud-sql-mssql.yaml │ │ ├── cloud-sql-mysql-admin.yaml │ │ ├── cloud-sql-mysql-observability.yaml │ │ ├── cloud-sql-mysql.yaml │ │ ├── cloud-sql-postgres-admin.yaml │ │ ├── cloud-sql-postgres-observability.yaml │ │ ├── cloud-sql-postgres.yaml │ │ ├── dataplex.yaml │ │ ├── firestore.yaml │ │ ├── looker-conversational-analytics.yaml │ │ ├── looker.yaml │ │ ├── mssql.yaml │ │ ├── mysql.yaml │ │ ├── neo4j.yaml │ │ ├── oceanbase.yaml │ │ ├── postgres.yaml │ │ ├── spanner-postgres.yaml │ │ ├── spanner.yaml │ │ └── sqlite.yaml │ ├── server │ │ ├── api_test.go │ │ ├── api.go │ │ ├── common_test.go │ │ ├── config.go │ │ ├── mcp │ │ │ ├── jsonrpc │ │ │ │ ├── jsonrpc_test.go │ │ │ │ └── jsonrpc.go │ │ │ ├── mcp.go │ │ │ ├── util │ │ │ │ └── lifecycle.go │ │ │ ├── v20241105 │ │ │ │ ├── method.go │ │ │ │ └── types.go │ │ │ ├── v20250326 │ │ │ │ ├── method.go │ │ │ │ └── types.go │ │ │ └── v20250618 │ │ │ ├── method.go │ │ │ └── types.go │ │ ├── mcp_test.go │ │ ├── mcp.go │ │ ├── server_test.go │ │ ├── server.go │ │ ├── static │ │ │ ├── assets │ │ │ │ └── mcptoolboxlogo.png │ │ │ ├── css │ │ │ │ └── style.css │ │ │ ├── index.html │ │ │ ├── js │ │ │ │ ├── auth.js │ │ │ │ ├── loadTools.js │ │ │ │ ├── mainContent.js │ │ │ │ ├── navbar.js │ │ │ │ ├── runTool.js │ │ │ │ ├── toolDisplay.js │ │ │ │ ├── tools.js │ │ │ │ └── toolsets.js │ │ │ ├── tools.html │ │ │ └── toolsets.html │ │ ├── web_test.go │ │ └── web.go │ ├── sources │ │ ├── alloydbadmin │ │ │ ├── alloydbadmin_test.go │ │ │ └── alloydbadmin.go │ │ ├── alloydbpg │ │ │ ├── alloydb_pg_test.go │ │ │ └── alloydb_pg.go │ │ ├── bigquery │ │ │ ├── bigquery_test.go │ │ │ └── bigquery.go │ │ ├── bigtable │ │ │ ├── bigtable_test.go │ │ │ └── bigtable.go │ │ ├── cassandra │ │ │ ├── cassandra_test.go │ │ │ └── cassandra.go │ │ ├── clickhouse │ │ │ ├── clickhouse_test.go │ │ │ └── clickhouse.go │ │ ├── cloudmonitoring │ │ │ ├── cloud_monitoring_test.go │ │ │ └── cloud_monitoring.go │ │ ├── cloudsqladmin │ │ │ ├── cloud_sql_admin_test.go │ │ │ └── cloud_sql_admin.go │ │ ├── cloudsqlmssql │ │ │ ├── cloud_sql_mssql_test.go │ │ │ └── cloud_sql_mssql.go │ │ ├── cloudsqlmysql │ │ │ ├── cloud_sql_mysql_test.go │ │ │ └── cloud_sql_mysql.go │ │ ├── cloudsqlpg │ │ │ ├── cloud_sql_pg_test.go │ │ │ └── cloud_sql_pg.go │ │ ├── couchbase │ │ │ ├── couchbase_test.go │ │ │ └── couchbase.go │ │ ├── dataplex │ │ │ ├── dataplex_test.go │ │ │ └── dataplex.go │ │ ├── dgraph │ │ │ ├── dgraph_test.go │ │ │ └── dgraph.go │ │ ├── dialect.go │ │ ├── firebird │ │ │ ├── firebird_test.go │ │ │ └── firebird.go │ │ ├── firestore │ │ │ ├── firestore_test.go │ │ │ └── firestore.go │ │ ├── http │ │ │ ├── http_test.go │ │ │ └── http.go │ │ ├── ip_type.go │ │ ├── looker │ │ │ ├── looker_test.go │ │ │ └── looker.go │ │ ├── mongodb │ │ │ ├── mongodb_test.go │ │ │ └── mongodb.go │ │ ├── mssql │ │ │ ├── mssql_test.go │ │ │ └── mssql.go │ │ ├── mysql │ │ │ ├── mysql_test.go │ │ │ └── mysql.go │ │ ├── neo4j │ │ │ ├── neo4j_test.go │ │ │ └── neo4j.go │ │ ├── oceanbase │ │ │ ├── oceanbase_test.go │ │ │ └── oceanbase.go │ │ ├── oracle │ │ │ └── oracle.go │ │ ├── postgres │ │ │ ├── postgres_test.go │ │ │ └── postgres.go │ │ ├── redis │ │ │ ├── redis_test.go │ │ │ └── redis.go │ │ ├── sources.go │ │ ├── spanner │ │ │ ├── spanner_test.go │ │ │ └── spanner.go │ │ ├── sqlite │ │ │ ├── sqlite_test.go │ │ │ └── sqlite.go │ │ ├── tidb │ │ │ ├── tidb_test.go │ │ │ └── tidb.go │ │ ├── trino │ │ │ ├── trino_test.go │ │ │ └── trino.go │ │ ├── util.go │ │ ├── valkey │ │ │ ├── valkey_test.go │ │ │ └── valkey.go │ │ └── yugabytedb │ │ ├── yugabytedb_test.go │ │ └── yugabytedb.go │ ├── telemetry │ │ ├── instrumentation.go │ │ └── telemetry.go │ ├── testutils │ │ └── testutils.go │ ├── tools │ │ ├── alloydb │ │ │ ├── alloydbcreatecluster │ │ │ │ ├── alloydbcreatecluster_test.go │ │ │ │ └── alloydbcreatecluster.go │ │ │ ├── alloydbcreateinstance │ │ │ │ ├── alloydbcreateinstance_test.go │ │ │ │ └── alloydbcreateinstance.go │ │ │ ├── alloydbcreateuser │ │ │ │ ├── alloydbcreateuser_test.go │ │ │ │ └── alloydbcreateuser.go │ │ │ ├── alloydbgetcluster │ │ │ │ ├── alloydbgetcluster_test.go │ │ │ │ └── alloydbgetcluster.go │ │ │ ├── alloydbgetinstance │ │ │ │ ├── alloydbgetinstance_test.go │ │ │ │ └── alloydbgetinstance.go │ │ │ ├── alloydbgetuser │ │ │ │ ├── alloydbgetuser_test.go │ │ │ │ └── alloydbgetuser.go │ │ │ ├── alloydblistclusters │ │ │ │ ├── alloydblistclusters_test.go │ │ │ │ └── alloydblistclusters.go │ │ │ ├── alloydblistinstances │ │ │ │ ├── alloydblistinstances_test.go │ │ │ │ └── alloydblistinstances.go │ │ │ ├── alloydblistusers │ │ │ │ ├── alloydblistusers_test.go │ │ │ │ └── alloydblistusers.go │ │ │ └── alloydbwaitforoperation │ │ │ ├── alloydbwaitforoperation_test.go │ │ │ └── alloydbwaitforoperation.go │ │ ├── alloydbainl │ │ │ ├── alloydbainl_test.go │ │ │ └── alloydbainl.go │ │ ├── bigquery │ │ │ ├── bigqueryanalyzecontribution │ │ │ │ ├── bigqueryanalyzecontribution_test.go │ │ │ │ └── bigqueryanalyzecontribution.go │ │ │ ├── bigquerycommon │ │ │ │ ├── table_name_parser_test.go │ │ │ │ ├── table_name_parser.go │ │ │ │ └── util.go │ │ │ ├── bigqueryconversationalanalytics │ │ │ │ ├── bigqueryconversationalanalytics_test.go │ │ │ │ └── bigqueryconversationalanalytics.go │ │ │ ├── bigqueryexecutesql │ │ │ │ ├── bigqueryexecutesql_test.go │ │ │ │ └── bigqueryexecutesql.go │ │ │ ├── bigqueryforecast │ │ │ │ ├── bigqueryforecast_test.go │ │ │ │ └── bigqueryforecast.go │ │ │ ├── bigquerygetdatasetinfo │ │ │ │ ├── bigquerygetdatasetinfo_test.go │ │ │ │ └── bigquerygetdatasetinfo.go │ │ │ ├── bigquerygettableinfo │ │ │ │ ├── bigquerygettableinfo_test.go │ │ │ │ └── bigquerygettableinfo.go │ │ │ ├── bigquerylistdatasetids │ │ │ │ ├── bigquerylistdatasetids_test.go │ │ │ │ └── bigquerylistdatasetids.go │ │ │ ├── bigquerylisttableids │ │ │ │ ├── bigquerylisttableids_test.go │ │ │ │ └── bigquerylisttableids.go │ │ │ ├── bigquerysearchcatalog │ │ │ │ ├── bigquerysearchcatalog_test.go │ │ │ │ └── bigquerysearchcatalog.go │ │ │ └── bigquerysql │ │ │ ├── bigquerysql_test.go │ │ │ └── bigquerysql.go │ │ ├── bigtable │ │ │ ├── bigtable_test.go │ │ │ └── bigtable.go │ │ ├── cassandra │ │ │ └── cassandracql │ │ │ ├── cassandracql_test.go │ │ │ └── cassandracql.go │ │ ├── clickhouse │ │ │ ├── clickhouseexecutesql │ │ │ │ ├── clickhouseexecutesql_test.go │ │ │ │ └── clickhouseexecutesql.go │ │ │ ├── clickhouselistdatabases │ │ │ │ ├── clickhouselistdatabases_test.go │ │ │ │ └── clickhouselistdatabases.go │ │ │ ├── clickhouselisttables │ │ │ │ ├── clickhouselisttables_test.go │ │ │ │ └── clickhouselisttables.go │ │ │ └── clickhousesql │ │ │ ├── clickhousesql_test.go │ │ │ └── clickhousesql.go │ │ ├── cloudmonitoring │ │ │ ├── cloudmonitoring_test.go │ │ │ └── cloudmonitoring.go │ │ ├── cloudsql │ │ │ ├── cloudsqlcreatedatabase │ │ │ │ ├── cloudsqlcreatedatabase_test.go │ │ │ │ └── cloudsqlcreatedatabase.go │ │ │ ├── cloudsqlcreateusers │ │ │ │ ├── cloudsqlcreateusers_test.go │ │ │ │ └── cloudsqlcreateusers.go │ │ │ ├── cloudsqlgetinstances │ │ │ │ ├── cloudsqlgetinstances_test.go │ │ │ │ └── cloudsqlgetinstances.go │ │ │ ├── cloudsqllistdatabases │ │ │ │ ├── cloudsqllistdatabases_test.go │ │ │ │ └── cloudsqllistdatabases.go │ │ │ ├── cloudsqllistinstances │ │ │ │ ├── cloudsqllistinstances_test.go │ │ │ │ └── cloudsqllistinstances.go │ │ │ └── cloudsqlwaitforoperation │ │ │ ├── cloudsqlwaitforoperation_test.go │ │ │ └── cloudsqlwaitforoperation.go │ │ ├── cloudsqlmssql │ │ │ └── cloudsqlmssqlcreateinstance │ │ │ ├── cloudsqlmssqlcreateinstance_test.go │ │ │ └── cloudsqlmssqlcreateinstance.go │ │ ├── cloudsqlmysql │ │ │ └── cloudsqlmysqlcreateinstance │ │ │ ├── cloudsqlmysqlcreateinstance_test.go │ │ │ └── cloudsqlmysqlcreateinstance.go │ │ ├── cloudsqlpg │ │ │ └── cloudsqlpgcreateinstances │ │ │ ├── cloudsqlpgcreateinstances_test.go │ │ │ └── cloudsqlpgcreateinstances.go │ │ ├── common_test.go │ │ ├── common.go │ │ ├── couchbase │ │ │ ├── couchbase_test.go │ │ │ └── couchbase.go │ │ ├── dataform │ │ │ └── dataformcompilelocal │ │ │ ├── dataformcompilelocal_test.go │ │ │ └── dataformcompilelocal.go │ │ ├── dataplex │ │ │ ├── dataplexlookupentry │ │ │ │ ├── dataplexlookupentry_test.go │ │ │ │ └── dataplexlookupentry.go │ │ │ ├── dataplexsearchaspecttypes │ │ │ │ ├── dataplexsearchaspecttypes_test.go │ │ │ │ └── dataplexsearchaspecttypes.go │ │ │ └── dataplexsearchentries │ │ │ ├── dataplexsearchentries_test.go │ │ │ └── dataplexsearchentries.go │ │ ├── dgraph │ │ │ ├── dgraph_test.go │ │ │ └── dgraph.go │ │ ├── firebird │ │ │ ├── firebirdexecutesql │ │ │ │ ├── firebirdexecutesql_test.go │ │ │ │ └── firebirdexecutesql.go │ │ │ └── firebirdsql │ │ │ ├── firebirdsql_test.go │ │ │ └── firebirdsql.go │ │ ├── firestore │ │ │ ├── firestoreadddocuments │ │ │ │ ├── firestoreadddocuments_test.go │ │ │ │ └── firestoreadddocuments.go │ │ │ ├── firestoredeletedocuments │ │ │ │ ├── firestoredeletedocuments_test.go │ │ │ │ └── firestoredeletedocuments.go │ │ │ ├── firestoregetdocuments │ │ │ │ ├── firestoregetdocuments_test.go │ │ │ │ └── firestoregetdocuments.go │ │ │ ├── firestoregetrules │ │ │ │ ├── firestoregetrules_test.go │ │ │ │ └── firestoregetrules.go │ │ │ ├── firestorelistcollections │ │ │ │ ├── firestorelistcollections_test.go │ │ │ │ └── firestorelistcollections.go │ │ │ ├── firestorequery │ │ │ │ ├── firestorequery_test.go │ │ │ │ └── firestorequery.go │ │ │ ├── firestorequerycollection │ │ │ │ ├── firestorequerycollection_test.go │ │ │ │ └── firestorequerycollection.go │ │ │ ├── firestoreupdatedocument │ │ │ │ ├── firestoreupdatedocument_test.go │ │ │ │ └── firestoreupdatedocument.go │ │ │ ├── firestorevalidaterules │ │ │ │ ├── firestorevalidaterules_test.go │ │ │ │ └── firestorevalidaterules.go │ │ │ └── util │ │ │ ├── converter_test.go │ │ │ ├── converter.go │ │ │ ├── validator_test.go │ │ │ └── validator.go │ │ ├── http │ │ │ ├── http_test.go │ │ │ └── http.go │ │ ├── http_method.go │ │ ├── looker │ │ │ ├── lookeradddashboardelement │ │ │ │ ├── lookeradddashboardelement_test.go │ │ │ │ └── lookeradddashboardelement.go │ │ │ ├── lookercommon │ │ │ │ ├── lookercommon_test.go │ │ │ │ └── lookercommon.go │ │ │ ├── lookerconversationalanalytics │ │ │ │ ├── lookerconversationalanalytics_test.go │ │ │ │ └── lookerconversationalanalytics.go │ │ │ ├── lookergetdashboards │ │ │ │ ├── lookergetdashboards_test.go │ │ │ │ └── lookergetdashboards.go │ │ │ ├── lookergetdimensions │ │ │ │ ├── lookergetdimensions_test.go │ │ │ │ └── lookergetdimensions.go │ │ │ ├── lookergetexplores │ │ │ │ ├── lookergetexplores_test.go │ │ │ │ └── lookergetexplores.go │ │ │ ├── lookergetfilters │ │ │ │ ├── lookergetfilters_test.go │ │ │ │ └── lookergetfilters.go │ │ │ ├── lookergetlooks │ │ │ │ ├── lookergetlooks_test.go │ │ │ │ └── lookergetlooks.go │ │ │ ├── lookergetmeasures │ │ │ │ ├── lookergetmeasures_test.go │ │ │ │ └── lookergetmeasures.go │ │ │ ├── lookergetmodels │ │ │ │ ├── lookergetmodels_test.go │ │ │ │ └── lookergetmodels.go │ │ │ ├── lookergetparameters │ │ │ │ ├── lookergetparameters_test.go │ │ │ │ └── lookergetparameters.go │ │ │ ├── lookerhealthanalyze │ │ │ │ ├── lookerhealthanalyze_test.go │ │ │ │ └── lookerhealthanalyze.go │ │ │ ├── lookerhealthpulse │ │ │ │ ├── lookerhealthpulse_test.go │ │ │ │ └── lookerhealthpulse.go │ │ │ ├── lookerhealthvacuum │ │ │ │ ├── lookerhealthvacuum_test.go │ │ │ │ └── lookerhealthvacuum.go │ │ │ ├── lookermakedashboard │ │ │ │ ├── lookermakedashboard_test.go │ │ │ │ └── lookermakedashboard.go │ │ │ ├── lookermakelook │ │ │ │ ├── lookermakelook_test.go │ │ │ │ └── lookermakelook.go │ │ │ ├── lookerquery │ │ │ │ ├── lookerquery_test.go │ │ │ │ └── lookerquery.go │ │ │ ├── lookerquerysql │ │ │ │ ├── lookerquerysql_test.go │ │ │ │ └── lookerquerysql.go │ │ │ ├── lookerqueryurl │ │ │ │ ├── lookerqueryurl_test.go │ │ │ │ └── lookerqueryurl.go │ │ │ └── lookerrunlook │ │ │ ├── lookerrunlook_test.go │ │ │ └── lookerrunlook.go │ │ ├── mongodb │ │ │ ├── mongodbaggregate │ │ │ │ ├── mongodbaggregate_test.go │ │ │ │ └── mongodbaggregate.go │ │ │ ├── mongodbdeletemany │ │ │ │ ├── mongodbdeletemany_test.go │ │ │ │ └── mongodbdeletemany.go │ │ │ ├── mongodbdeleteone │ │ │ │ ├── mongodbdeleteone_test.go │ │ │ │ └── mongodbdeleteone.go │ │ │ ├── mongodbfind │ │ │ │ ├── mongodbfind_test.go │ │ │ │ └── mongodbfind.go │ │ │ ├── mongodbfindone │ │ │ │ ├── mongodbfindone_test.go │ │ │ │ └── mongodbfindone.go │ │ │ ├── mongodbinsertmany │ │ │ │ ├── mongodbinsertmany_test.go │ │ │ │ └── mongodbinsertmany.go │ │ │ ├── mongodbinsertone │ │ │ │ ├── mongodbinsertone_test.go │ │ │ │ └── mongodbinsertone.go │ │ │ ├── mongodbupdatemany │ │ │ │ ├── mongodbupdatemany_test.go │ │ │ │ └── mongodbupdatemany.go │ │ │ └── mongodbupdateone │ │ │ ├── mongodbupdateone_test.go │ │ │ └── mongodbupdateone.go │ │ ├── mssql │ │ │ ├── mssqlexecutesql │ │ │ │ ├── mssqlexecutesql_test.go │ │ │ │ └── mssqlexecutesql.go │ │ │ ├── mssqllisttables │ │ │ │ ├── mssqllisttables_test.go │ │ │ │ └── mssqllisttables.go │ │ │ └── mssqlsql │ │ │ ├── mssqlsql_test.go │ │ │ └── mssqlsql.go │ │ ├── mysql │ │ │ ├── mysqlcommon │ │ │ │ └── mysqlcommon.go │ │ │ ├── mysqlexecutesql │ │ │ │ ├── mysqlexecutesql_test.go │ │ │ │ └── mysqlexecutesql.go │ │ │ ├── mysqllistactivequeries │ │ │ │ ├── mysqllistactivequeries_test.go │ │ │ │ └── mysqllistactivequeries.go │ │ │ ├── mysqllisttablefragmentation │ │ │ │ ├── mysqllisttablefragmentation_test.go │ │ │ │ └── mysqllisttablefragmentation.go │ │ │ ├── mysqllisttables │ │ │ │ ├── mysqllisttables_test.go │ │ │ │ └── mysqllisttables.go │ │ │ ├── mysqllisttablesmissinguniqueindexes │ │ │ │ ├── mysqllisttablesmissinguniqueindexes_test.go │ │ │ │ └── mysqllisttablesmissinguniqueindexes.go │ │ │ └── mysqlsql │ │ │ ├── mysqlsql_test.go │ │ │ └── mysqlsql.go │ │ ├── neo4j │ │ │ ├── neo4jcypher │ │ │ │ ├── neo4jcypher_test.go │ │ │ │ └── neo4jcypher.go │ │ │ ├── neo4jexecutecypher │ │ │ │ ├── classifier │ │ │ │ │ ├── classifier_test.go │ │ │ │ │ └── classifier.go │ │ │ │ ├── neo4jexecutecypher_test.go │ │ │ │ └── neo4jexecutecypher.go │ │ │ └── neo4jschema │ │ │ ├── cache │ │ │ │ ├── cache_test.go │ │ │ │ └── cache.go │ │ │ ├── helpers │ │ │ │ ├── helpers_test.go │ │ │ │ └── helpers.go │ │ │ ├── neo4jschema_test.go │ │ │ ├── neo4jschema.go │ │ │ └── types │ │ │ └── types.go │ │ ├── oceanbase │ │ │ ├── oceanbaseexecutesql │ │ │ │ ├── oceanbaseexecutesql_test.go │ │ │ │ └── oceanbaseexecutesql.go │ │ │ └── oceanbasesql │ │ │ ├── oceanbasesql_test.go │ │ │ └── oceanbasesql.go │ │ ├── oracle │ │ │ ├── oracleexecutesql │ │ │ │ └── oracleexecutesql.go │ │ │ └── oraclesql │ │ │ └── oraclesql.go │ │ ├── parameters_test.go │ │ ├── parameters.go │ │ ├── postgres │ │ │ ├── postgresexecutesql │ │ │ │ ├── postgresexecutesql_test.go │ │ │ │ └── postgresexecutesql.go │ │ │ ├── postgreslistactivequeries │ │ │ │ ├── postgreslistactivequeries_test.go │ │ │ │ └── postgreslistactivequeries.go │ │ │ ├── postgreslistavailableextensions │ │ │ │ ├── postgreslistavailableextensions_test.go │ │ │ │ └── postgreslistavailableextensions.go │ │ │ ├── postgreslistinstalledextensions │ │ │ │ ├── postgreslistinstalledextensions_test.go │ │ │ │ └── postgreslistinstalledextensions.go │ │ │ ├── postgreslisttables │ │ │ │ ├── postgreslisttables_test.go │ │ │ │ └── postgreslisttables.go │ │ │ └── postgressql │ │ │ ├── postgressql_test.go │ │ │ └── postgressql.go │ │ ├── redis │ │ │ ├── redis_test.go │ │ │ └── redis.go │ │ ├── spanner │ │ │ ├── spannerexecutesql │ │ │ │ ├── spannerexecutesql_test.go │ │ │ │ └── spannerexecutesql.go │ │ │ ├── spannerlisttables │ │ │ │ ├── spannerlisttables_test.go │ │ │ │ └── spannerlisttables.go │ │ │ └── spannersql │ │ │ ├── spanner_test.go │ │ │ └── spannersql.go │ │ ├── sqlite │ │ │ ├── sqliteexecutesql │ │ │ │ ├── sqliteexecutesql_test.go │ │ │ │ └── sqliteexecutesql.go │ │ │ └── sqlitesql │ │ │ ├── sqlitesql_test.go │ │ │ └── sqlitesql.go │ │ ├── tidb │ │ │ ├── tidbexecutesql │ │ │ │ ├── tidbexecutesql_test.go │ │ │ │ └── tidbexecutesql.go │ │ │ └── tidbsql │ │ │ ├── tidbsql_test.go │ │ │ └── tidbsql.go │ │ ├── tools_test.go │ │ ├── tools.go │ │ ├── toolsets.go │ │ ├── trino │ │ │ ├── trinoexecutesql │ │ │ │ ├── trinoexecutesql_test.go │ │ │ │ └── trinoexecutesql.go │ │ │ └── trinosql │ │ │ ├── trinosql_test.go │ │ │ └── trinosql.go │ │ ├── utility │ │ │ └── wait │ │ │ ├── wait_test.go │ │ │ └── wait.go │ │ ├── valkey │ │ │ ├── valkey_test.go │ │ │ └── valkey.go │ │ └── yugabytedbsql │ │ ├── yugabytedbsql_test.go │ │ └── yugabytedbsql.go │ └── util │ └── util.go ├── LICENSE ├── logo.png ├── main.go ├── README.md └── tests ├── alloydb │ ├── alloydb_integration_test.go │ └── alloydb_wait_for_operation_test.go ├── alloydbainl │ └── alloydb_ai_nl_integration_test.go ├── alloydbpg │ └── alloydb_pg_integration_test.go ├── auth.go ├── bigquery │ └── bigquery_integration_test.go ├── bigtable │ └── bigtable_integration_test.go ├── cassandra │ └── cassandra_integration_test.go ├── clickhouse │ └── clickhouse_integration_test.go ├── cloudmonitoring │ └── cloud_monitoring_integration_test.go ├── cloudsql │ ├── cloud_sql_create_database_test.go │ ├── cloud_sql_create_users_test.go │ ├── cloud_sql_get_instances_test.go │ ├── cloud_sql_list_databases_test.go │ ├── cloudsql_list_instances_test.go │ └── cloudsql_wait_for_operation_test.go ├── cloudsqlmssql │ ├── cloud_sql_mssql_create_instance_integration_test.go │ └── cloud_sql_mssql_integration_test.go ├── cloudsqlmysql │ ├── cloud_sql_mysql_create_instance_integration_test.go │ └── cloud_sql_mysql_integration_test.go ├── cloudsqlpg │ ├── cloud_sql_pg_create_instances_test.go │ └── cloud_sql_pg_integration_test.go ├── common.go ├── couchbase │ └── couchbase_integration_test.go ├── dataform │ └── dataform_integration_test.go ├── dataplex │ └── dataplex_integration_test.go ├── dgraph │ └── dgraph_integration_test.go ├── firebird │ └── firebird_integration_test.go ├── firestore │ └── firestore_integration_test.go ├── http │ └── http_integration_test.go ├── looker │ └── looker_integration_test.go ├── mongodb │ └── mongodb_integration_test.go ├── mssql │ └── mssql_integration_test.go ├── mysql │ └── mysql_integration_test.go ├── neo4j │ └── neo4j_integration_test.go ├── oceanbase │ └── oceanbase_integration_test.go ├── option.go ├── oracle │ └── oracle_integration_test.go ├── postgres │ └── postgres_integration_test.go ├── redis │ └── redis_test.go ├── server.go ├── source.go ├── spanner │ └── spanner_integration_test.go ├── sqlite │ └── sqlite_integration_test.go ├── tidb │ └── tidb_integration_test.go ├── tool.go ├── trino │ └── trino_integration_test.go ├── utility │ └── wait_integration_test.go ├── valkey │ └── valkey_test.go └── yugabytedb └── yugabytedb_integration_test.go ``` # Files -------------------------------------------------------------------------------- /tests/firestore/firestore_integration_test.go: -------------------------------------------------------------------------------- ```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 firestore import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "os" "reflect" "regexp" "strings" "testing" "time" firestoreapi "cloud.google.com/go/firestore" "github.com/google/uuid" "github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc" "github.com/googleapis/genai-toolbox/internal/testutils" "github.com/googleapis/genai-toolbox/tests" "google.golang.org/api/option" ) var ( FirestoreSourceKind = "firestore" FirestoreProject = os.Getenv("FIRESTORE_PROJECT") FirestoreDatabase = os.Getenv("FIRESTORE_DATABASE") // Optional, defaults to "(default)" ) func getFirestoreVars(t *testing.T) map[string]any { if FirestoreProject == "" { t.Fatal("'FIRESTORE_PROJECT' not set") } vars := map[string]any{ "kind": FirestoreSourceKind, "project": FirestoreProject, } // Only add database if it's explicitly set if FirestoreDatabase != "" { vars["database"] = FirestoreDatabase } return vars } // initFirestoreConnection creates a Firestore client for testing func initFirestoreConnection(project, database string) (*firestoreapi.Client, error) { ctx := context.Background() if database == "" { database = "(default)" } client, err := firestoreapi.NewClientWithDatabase(ctx, project, database, option.WithUserAgent("genai-toolbox-integration-test")) if err != nil { return nil, fmt.Errorf("failed to create Firestore client for project %q and database %q: %w", project, database, err) } return client, nil } func TestFirestoreToolEndpoints(t *testing.T) { sourceConfig := getFirestoreVars(t) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) defer cancel() var args []string client, err := initFirestoreConnection(FirestoreProject, FirestoreDatabase) if err != nil { t.Fatalf("unable to create Firestore connection: %s", err) } defer client.Close() // Create test collection and document names with UUID testCollectionName := fmt.Sprintf("test_collection_%s", strings.ReplaceAll(uuid.New().String(), "-", "")) testSubCollectionName := fmt.Sprintf("test_subcollection_%s", strings.ReplaceAll(uuid.New().String(), "-", "")) testDocID1 := fmt.Sprintf("doc_%s", strings.ReplaceAll(uuid.New().String(), "-", "")) testDocID2 := fmt.Sprintf("doc_%s", strings.ReplaceAll(uuid.New().String(), "-", "")) testDocID3 := fmt.Sprintf("doc_%s", strings.ReplaceAll(uuid.New().String(), "-", "")) // Document paths for testing docPath1 := fmt.Sprintf("%s/%s", testCollectionName, testDocID1) docPath2 := fmt.Sprintf("%s/%s", testCollectionName, testDocID2) docPath3 := fmt.Sprintf("%s/%s", testCollectionName, testDocID3) // Set up test data teardown := setupFirestoreTestData(t, ctx, client, testCollectionName, testSubCollectionName, testDocID1, testDocID2, testDocID3) defer teardown(t) // Write config into a file and pass it to command toolsFile := getFirestoreToolsConfig(sourceConfig) cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...) if err != nil { t.Fatalf("command initialization returned an error: %s", err) } defer cleanup() waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second) defer cancel() 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: %s", err) } // Run Firestore-specific tool get test runFirestoreToolGetTest(t) // Run Firestore-specific MCP test runFirestoreMCPToolCallMethod(t, docPath1, docPath2) // Run specific Firestore tool tests runFirestoreGetDocumentsTest(t, docPath1, docPath2) runFirestoreQueryCollectionTest(t, testCollectionName) runFirestoreQueryTest(t, testCollectionName) runFirestoreQuerySelectArrayTest(t, testCollectionName) runFirestoreListCollectionsTest(t, testCollectionName, testSubCollectionName, docPath1) runFirestoreAddDocumentsTest(t, testCollectionName) runFirestoreUpdateDocumentTest(t, testCollectionName, testDocID1) runFirestoreDeleteDocumentsTest(t, docPath3) runFirestoreGetRulesTest(t) runFirestoreValidateRulesTest(t) } func runFirestoreToolGetTest(t *testing.T) { // Test tool get endpoint for Firestore tools tcs := []struct { name string api string want map[string]any }{ { name: "get my-simple-tool", api: "http://127.0.0.1:5000/api/tool/my-simple-tool/", want: map[string]any{ "my-simple-tool": map[string]any{ "description": "Simple tool to test end to end functionality.", "parameters": []any{ map[string]any{ "name": "documentPaths", "type": "array", "required": true, "description": "Array of document paths to retrieve from Firestore.", "items": map[string]any{ "name": "item", "type": "string", "required": true, "description": "Document path", "authSources": []any{}, }, "authSources": []any{}, }, }, "authRequired": []any{}, }, }, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { resp, err := http.Get(tc.api) if err != nil { t.Fatalf("error when sending a request: %s", err) } defer resp.Body.Close() if resp.StatusCode != 200 { t.Fatalf("response status code is not 200") } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body") } got, ok := body["tools"] if !ok { t.Fatalf("unable to find tools in response body") } // Compare as JSON strings to handle any ordering differences gotJSON, _ := json.Marshal(got) wantJSON, _ := json.Marshal(tc.want) if string(gotJSON) != string(wantJSON) { t.Logf("got %v, want %v", got, tc.want) } }) } } func runFirestoreValidateRulesTest(t *testing.T) { invokeTcs := []struct { name string api string requestBody io.Reader wantRegex string isErr bool }{ { name: "validate valid rules", api: "http://127.0.0.1:5000/api/tool/firestore-validate-rules/invoke", requestBody: bytes.NewBuffer([]byte(`{ "source": "rules_version = '2';\nservice cloud.firestore {\n match /databases/{database}/documents {\n match /{document=**} {\n allow read, write: if true;\n }\n }\n}" }`)), wantRegex: `"valid":true.*"issueCount":0`, isErr: false, }, { name: "validate rules with syntax error", api: "http://127.0.0.1:5000/api/tool/firestore-validate-rules/invoke", requestBody: bytes.NewBuffer([]byte(`{ "source": "rules_version = '2';\nservice cloud.firestore {\n match /databases/{database}/documents {\n match /{document=**} {\n allow read, write: if true;;\n }\n }\n}" }`)), wantRegex: `"valid":false.*"issueCount":[1-9]`, isErr: false, }, { name: "validate rules with missing version", api: "http://127.0.0.1:5000/api/tool/firestore-validate-rules/invoke", requestBody: bytes.NewBuffer([]byte(`{ "source": "service cloud.firestore {\n match /databases/{database}/documents {\n match /{document=**} {\n allow read, write: if true;\n }\n }\n}" }`)), wantRegex: `"valid":false.*"issueCount":[1-9]`, isErr: false, }, { name: "validate empty rules", api: "http://127.0.0.1:5000/api/tool/firestore-validate-rules/invoke", requestBody: bytes.NewBuffer([]byte(`{"source": ""}`)), isErr: true, }, { name: "missing source parameter", api: "http://127.0.0.1:5000/api/tool/firestore-validate-rules/invoke", requestBody: bytes.NewBuffer([]byte(`{}`)), isErr: true, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %v", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if tc.wantRegex != "" { matched, err := regexp.MatchString(tc.wantRegex, got) if err != nil { t.Fatalf("invalid regex pattern: %v", err) } if !matched { t.Fatalf("result does not match expected pattern.\nGot: %s\nWant pattern: %s", got, tc.wantRegex) } } }) } } func runFirestoreGetRulesTest(t *testing.T) { invokeTcs := []struct { name string api string requestBody io.Reader wantRegex string isErr bool }{ { name: "get firestore rules", api: "http://127.0.0.1:5000/api/tool/firestore-get-rules/invoke", requestBody: bytes.NewBuffer([]byte(`{}`)), wantRegex: `"content":"[^"]+"`, // Should contain at least one of these fields isErr: false, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { bodyBytes, _ := io.ReadAll(resp.Body) // The test might fail if there are no active rules in the project, which is acceptable if strings.Contains(string(bodyBytes), "no active Firestore rules") { t.Skipf("No active Firestore rules found in the project") return } if tc.isErr { return } t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %v", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if tc.wantRegex != "" { matched, err := regexp.MatchString(tc.wantRegex, got) if err != nil { t.Fatalf("invalid regex pattern: %v", err) } if !matched { t.Fatalf("result does not match expected pattern.\nGot: %s\nWant pattern: %s", got, tc.wantRegex) } } }) } } func runFirestoreMCPToolCallMethod(t *testing.T, docPath1, docPath2 string) { sessionId := tests.RunInitialize(t, "2024-11-05") header := map[string]string{} if sessionId != "" { header["Mcp-Session-Id"] = sessionId } // Test tool invoke endpoint invokeTcs := []struct { name string api string requestBody jsonrpc.JSONRPCRequest requestHeader map[string]string wantContains string wantError bool }{ { name: "MCP Invoke my-param-tool", api: "http://127.0.0.1:5000/mcp", requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "my-param-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-param-tool", "arguments": map[string]any{ "documentPaths": []string{docPath1}, }, }, }, wantContains: `\"name\":\"Alice\"`, wantError: false, }, { name: "MCP Invoke invalid tool", api: "http://127.0.0.1:5000/mcp", requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invalid-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "foo", "arguments": map[string]any{}, }, }, wantContains: `tool with name \"foo\" does not exist`, wantError: true, }, { name: "MCP Invoke my-param-tool without parameters", api: "http://127.0.0.1:5000/mcp", requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke-without-parameter", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-param-tool", "arguments": map[string]any{}, }, }, wantContains: `parameter \"documentPaths\" is required`, wantError: true, }, { name: "MCP Invoke my-auth-required-tool", api: "http://127.0.0.1:5000/mcp", requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke my-auth-required-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-auth-required-tool", "arguments": map[string]any{}, }, }, wantContains: `tool with name \"my-auth-required-tool\" does not exist`, wantError: true, }, { name: "MCP Invoke my-fail-tool", api: "http://127.0.0.1:5000/mcp", requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke-fail-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-fail-tool", "arguments": map[string]any{ "documentPaths": []string{"non-existent/path"}, }, }, }, wantContains: `\"exists\":false`, wantError: false, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { reqMarshal, err := json.Marshal(tc.requestBody) if err != nil { t.Fatalf("unexpected error during marshaling of request body") } req, err := http.NewRequest(http.MethodPost, tc.api, bytes.NewBuffer(reqMarshal)) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") for k, v := range header { req.Header.Add(k, v) } resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() respBody, err := io.ReadAll(resp.Body) if err != nil { t.Fatalf("unable to read request body: %s", err) } got := string(bytes.TrimSpace(respBody)) if !strings.Contains(got, tc.wantContains) { t.Fatalf("Expected substring not found:\ngot: %q\nwant: %q (to be contained within got)", got, tc.wantContains) } }) } } func getFirestoreToolsConfig(sourceConfig map[string]any) map[string]any { sources := map[string]any{ "my-instance": sourceConfig, } tools := map[string]any{ // Tool for RunToolGetTest "my-simple-tool": map[string]any{ "kind": "firestore-get-documents", "source": "my-instance", "description": "Simple tool to test end to end functionality.", }, // Tool for MCP test - this will get documents "my-param-tool": map[string]any{ "kind": "firestore-get-documents", "source": "my-instance", "description": "Tool to get documents by paths", }, // Tool for MCP test that fails "my-fail-tool": map[string]any{ "kind": "firestore-get-documents", "source": "my-instance", "description": "Tool that will fail", }, // Firestore specific tools "firestore-get-docs": map[string]any{ "kind": "firestore-get-documents", "source": "my-instance", "description": "Get multiple documents from Firestore", }, "firestore-list-colls": map[string]any{ "kind": "firestore-list-collections", "source": "my-instance", "description": "List Firestore collections", }, "firestore-delete-docs": map[string]any{ "kind": "firestore-delete-documents", "source": "my-instance", "description": "Delete documents from Firestore", }, "firestore-query-coll": map[string]any{ "kind": "firestore-query-collection", "source": "my-instance", "description": "Query a Firestore collection", }, "firestore-query-param": map[string]any{ "kind": "firestore-query", "source": "my-instance", "description": "Query a Firestore collection with parameterizable filters", "collectionPath": "{{.collection}}", "filters": `{ "field": "age", "op": "{{.operator}}", "value": {"integerValue": "{{.ageValue}}"} }`, "limit": 10, "parameters": []map[string]any{ { "name": "collection", "type": "string", "description": "Collection to query", "required": true, }, { "name": "operator", "type": "string", "description": "Comparison operator", "required": true, }, { "name": "ageValue", "type": "string", "description": "Age value to compare", "required": true, }, }, }, "firestore-query-select-array": map[string]any{ "kind": "firestore-query", "source": "my-instance", "description": "Query with array-based select fields", "collectionPath": "{{.collection}}", "select": []string{"{{.fields}}"}, "limit": 10, "parameters": []map[string]any{ { "name": "collection", "type": "string", "description": "Collection to query", "required": true, }, { "name": "fields", "type": "array", "description": "Fields to select", "required": true, "items": map[string]any{ "name": "field", "type": "string", "description": "field", }, }, }, }, "firestore-get-rules": map[string]any{ "kind": "firestore-get-rules", "source": "my-instance", "description": "Get Firestore security rules", }, "firestore-validate-rules": map[string]any{ "kind": "firestore-validate-rules", "source": "my-instance", "description": "Validate Firestore security rules", }, "firestore-add-docs": map[string]any{ "kind": "firestore-add-documents", "source": "my-instance", "description": "Add documents to Firestore", }, "firestore-update-doc": map[string]any{ "kind": "firestore-update-document", "source": "my-instance", "description": "Update a document in Firestore", }, } return map[string]any{ "sources": sources, "tools": tools, } } func runFirestoreUpdateDocumentTest(t *testing.T, collectionName string, docID string) { docPath := fmt.Sprintf("%s/%s", collectionName, docID) invokeTcs := []struct { name string api string requestBody io.Reader wantKeys []string validateContent bool expectedContent map[string]interface{} isErr bool }{ { name: "update document with simple fields", api: "http://127.0.0.1:5000/api/tool/firestore-update-doc/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "documentPath": "%s", "documentData": { "name": {"stringValue": "Alice Updated"}, "status": {"stringValue": "active"} } }`, docPath))), wantKeys: []string{"documentPath", "updateTime"}, isErr: false, }, { name: "update document with selective fields using updateMask", api: "http://127.0.0.1:5000/api/tool/firestore-update-doc/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "documentPath": "%s", "documentData": { "age": {"integerValue": "31"}, "email": {"stringValue": "[email protected]"} }, "updateMask": ["age"] }`, docPath))), wantKeys: []string{"documentPath", "updateTime"}, isErr: false, }, { name: "update document with field deletion", api: "http://127.0.0.1:5000/api/tool/firestore-update-doc/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "documentPath": "%s", "documentData": { "name": {"stringValue": "Alice Final"} }, "updateMask": ["name", "status"] }`, docPath))), wantKeys: []string{"documentPath", "updateTime"}, isErr: false, }, { name: "update document with complex types", api: "http://127.0.0.1:5000/api/tool/firestore-update-doc/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "documentPath": "%s", "documentData": { "location": { "geoPointValue": { "latitude": 40.7128, "longitude": -74.0060 } }, "tags": { "arrayValue": { "values": [ {"stringValue": "updated"}, {"stringValue": "test"} ] } }, "metadata": { "mapValue": { "fields": { "lastModified": {"timestampValue": "2025-01-15T10:00:00Z"}, "version": {"integerValue": "2"} } } } } }`, docPath))), wantKeys: []string{"documentPath", "updateTime"}, isErr: false, }, { name: "update document with returnData", api: "http://127.0.0.1:5000/api/tool/firestore-update-doc/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "documentPath": "%s", "documentData": { "testField": {"stringValue": "test value"}, "testNumber": {"integerValue": "42"} }, "returnData": true }`, docPath))), wantKeys: []string{"documentPath", "updateTime", "documentData"}, validateContent: true, expectedContent: map[string]interface{}{ "testField": "test value", "testNumber": float64(42), // JSON numbers are decoded as float64 }, isErr: false, }, { name: "update nested fields with updateMask", api: "http://127.0.0.1:5000/api/tool/firestore-update-doc/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "documentPath": "%s", "documentData": { "profile": { "mapValue": { "fields": { "bio": {"stringValue": "Updated bio"}, "avatar": {"stringValue": "avatar.jpg"} } } } }, "updateMask": ["profile.bio", "profile.avatar"] }`, docPath))), wantKeys: []string{"documentPath", "updateTime"}, isErr: false, }, { name: "missing documentPath parameter", api: "http://127.0.0.1:5000/api/tool/firestore-update-doc/invoke", requestBody: bytes.NewBuffer([]byte(`{"documentData": {"test": {"stringValue": "value"}}}`)), isErr: true, }, { name: "missing documentData parameter", api: "http://127.0.0.1:5000/api/tool/firestore-update-doc/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"documentPath": "%s"}`, docPath))), isErr: true, }, { name: "update non-existent document", api: "http://127.0.0.1:5000/api/tool/firestore-update-doc/invoke", requestBody: bytes.NewBuffer([]byte(`{ "documentPath": "non-existent-collection/non-existent-doc", "documentData": { "field": {"stringValue": "value"} } }`)), wantKeys: []string{"documentPath", "updateTime"}, // Set with MergeAll creates if doesn't exist isErr: false, }, { name: "invalid field in updateMask", api: "http://127.0.0.1:5000/api/tool/firestore-update-doc/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "documentPath": "%s", "documentData": { "field1": {"stringValue": "value1"} }, "updateMask": ["field1", "nonExistentField"] }`, docPath))), isErr: true, // Should fail because nonExistentField is not in documentData }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %v", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } // Parse the result string as JSON var resultJSON map[string]interface{} err = json.Unmarshal([]byte(got), &resultJSON) if err != nil { t.Fatalf("error parsing result as JSON: %v", err) } // Check if all wanted keys exist for _, key := range tc.wantKeys { if _, exists := resultJSON[key]; !exists { t.Fatalf("expected key %q not found in result: %s", key, got) } } // Validate document data if required if tc.validateContent { docData, ok := resultJSON["documentData"].(map[string]interface{}) if !ok { t.Fatalf("documentData is not a map: %v", resultJSON["documentData"]) } // Check that expected fields are present with correct values for key, expectedValue := range tc.expectedContent { actualValue, exists := docData[key] if !exists { t.Fatalf("expected field %q not found in documentData", key) } if actualValue != expectedValue { t.Fatalf("field %q mismatch: expected %v, got %v", key, expectedValue, actualValue) } } } }) } } func runFirestoreAddDocumentsTest(t *testing.T, collectionName string) { invokeTcs := []struct { name string api string requestBody io.Reader wantKeys []string validateDocData bool expectedDocData map[string]interface{} isErr bool }{ { name: "add document with simple types", api: "http://127.0.0.1:5000/api/tool/firestore-add-docs/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collectionPath": "%s", "documentData": { "name": {"stringValue": "Test User"}, "age": {"integerValue": "42"}, "score": {"doubleValue": 99.5}, "active": {"booleanValue": true}, "notes": {"nullValue": null} } }`, collectionName))), wantKeys: []string{"documentPath", "createTime"}, isErr: false, }, { name: "add document with complex types", api: "http://127.0.0.1:5000/api/tool/firestore-add-docs/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collectionPath": "%s", "documentData": { "location": { "geoPointValue": { "latitude": 37.7749, "longitude": -122.4194 } }, "timestamp": { "timestampValue": "2025-01-07T10:00:00Z" }, "tags": { "arrayValue": { "values": [ {"stringValue": "tag1"}, {"stringValue": "tag2"} ] } }, "metadata": { "mapValue": { "fields": { "version": {"integerValue": "1"}, "type": {"stringValue": "test"} } } } } }`, collectionName))), wantKeys: []string{"documentPath", "createTime"}, isErr: false, }, { name: "add document with returnData", api: "http://127.0.0.1:5000/api/tool/firestore-add-docs/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collectionPath": "%s", "documentData": { "name": {"stringValue": "Return Test"}, "value": {"integerValue": "123"} }, "returnData": true }`, collectionName))), wantKeys: []string{"documentPath", "createTime", "documentData"}, validateDocData: true, expectedDocData: map[string]interface{}{ "name": "Return Test", "value": float64(123), // JSON numbers are decoded as float64 }, isErr: false, }, { name: "add document with nested maps and arrays", api: "http://127.0.0.1:5000/api/tool/firestore-add-docs/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collectionPath": "%s", "documentData": { "company": { "mapValue": { "fields": { "name": {"stringValue": "Tech Corp"}, "employees": { "arrayValue": { "values": [ { "mapValue": { "fields": { "name": {"stringValue": "John"}, "role": {"stringValue": "Developer"} } } }, { "mapValue": { "fields": { "name": {"stringValue": "Jane"}, "role": {"stringValue": "Manager"} } } } ] } } } } } } }`, collectionName))), wantKeys: []string{"documentPath", "createTime"}, isErr: false, }, { name: "missing collectionPath parameter", api: "http://127.0.0.1:5000/api/tool/firestore-add-docs/invoke", requestBody: bytes.NewBuffer([]byte(`{"documentData": {"test": {"stringValue": "value"}}}`)), isErr: true, }, { name: "missing documentData parameter", api: "http://127.0.0.1:5000/api/tool/firestore-add-docs/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"collectionPath": "%s"}`, collectionName))), isErr: true, }, { name: "invalid documentData format", api: "http://127.0.0.1:5000/api/tool/firestore-add-docs/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"collectionPath": "%s", "documentData": "not an object"}`, collectionName))), isErr: true, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %v", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } // Parse the result string as JSON var resultJSON map[string]interface{} err = json.Unmarshal([]byte(got), &resultJSON) if err != nil { t.Fatalf("error parsing result as JSON: %v", err) } // Check if all wanted keys exist for _, key := range tc.wantKeys { if _, exists := resultJSON[key]; !exists { t.Fatalf("expected key %q not found in result: %s", key, got) } } // Validate document data if required if tc.validateDocData { docData, ok := resultJSON["documentData"].(map[string]interface{}) if !ok { t.Fatalf("documentData is not a map: %v", resultJSON["documentData"]) } // Use reflect.DeepEqual to compare the document data if !reflect.DeepEqual(docData, tc.expectedDocData) { t.Fatalf("documentData mismatch:\nexpected: %v\nactual: %v", tc.expectedDocData, docData) } } }) } } func setupFirestoreTestData(t *testing.T, ctx context.Context, client *firestoreapi.Client, collectionName, subCollectionName, docID1, docID2, docID3 string) func(*testing.T) { // Create test documents testData1 := map[string]interface{}{ "name": "Alice", "age": 30, } testData2 := map[string]interface{}{ "name": "Bob", "age": 25, } testData3 := map[string]interface{}{ "name": "Charlie", "age": 35, } // Create documents _, err := client.Collection(collectionName).Doc(docID1).Set(ctx, testData1) if err != nil { t.Fatalf("Failed to create test document 1: %v", err) } _, err = client.Collection(collectionName).Doc(docID2).Set(ctx, testData2) if err != nil { t.Fatalf("Failed to create test document 2: %v", err) } _, err = client.Collection(collectionName).Doc(docID3).Set(ctx, testData3) if err != nil { t.Fatalf("Failed to create test document 3: %v", err) } // Create a subcollection document subDocData := map[string]interface{}{ "type": "subcollection_doc", "value": "test", } _, err = client.Collection(collectionName).Doc(docID1).Collection(subCollectionName).Doc("subdoc1").Set(ctx, subDocData) if err != nil { t.Fatalf("Failed to create subcollection document: %v", err) } // Return cleanup function that deletes ALL collections and documents in the database return func(t *testing.T) { // Helper function to recursively delete all documents in a collection var deleteCollection func(*firestoreapi.CollectionRef) error deleteCollection = func(collection *firestoreapi.CollectionRef) error { // Get all documents in the collection docs, err := collection.Documents(ctx).GetAll() if err != nil { return fmt.Errorf("failed to list documents in collection %s: %w", collection.Path, err) } // Delete each document and its subcollections for _, doc := range docs { // First, get all subcollections of this document subcollections, err := doc.Ref.Collections(ctx).GetAll() if err != nil { return fmt.Errorf("failed to list subcollections of document %s: %w", doc.Ref.Path, err) } // Recursively delete each subcollection for _, subcoll := range subcollections { if err := deleteCollection(subcoll); err != nil { return fmt.Errorf("failed to delete subcollection %s: %w", subcoll.Path, err) } } // Delete the document itself if _, err := doc.Ref.Delete(ctx); err != nil { return fmt.Errorf("failed to delete document %s: %w", doc.Ref.Path, err) } } return nil } // Get all root collections in the database rootCollections, err := client.Collections(ctx).GetAll() if err != nil { t.Errorf("Failed to list root collections: %v", err) return } // Delete each root collection and all its contents for _, collection := range rootCollections { if err := deleteCollection(collection); err != nil { t.Errorf("Failed to delete collection %s and its contents: %v", collection.ID, err) } } t.Logf("Successfully deleted all collections and documents in the database") } } func runFirestoreGetDocumentsTest(t *testing.T, docPath1, docPath2 string) { invokeTcs := []struct { name string api string requestBody io.Reader wantRegex string isErr bool }{ { name: "get single document", api: "http://127.0.0.1:5000/api/tool/firestore-get-docs/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"documentPaths": ["%s"]}`, docPath1))), wantRegex: `"name":"Alice"`, isErr: false, }, { name: "get multiple documents", api: "http://127.0.0.1:5000/api/tool/firestore-get-docs/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"documentPaths": ["%s", "%s"]}`, docPath1, docPath2))), wantRegex: `"name":"Alice".*"name":"Bob"`, isErr: false, }, { name: "get non-existent document", api: "http://127.0.0.1:5000/api/tool/firestore-get-docs/invoke", requestBody: bytes.NewBuffer([]byte(`{"documentPaths": ["non-existent-collection/non-existent-doc"]}`)), wantRegex: `"exists":false`, isErr: false, }, { name: "missing documentPaths parameter", api: "http://127.0.0.1:5000/api/tool/firestore-get-docs/invoke", requestBody: bytes.NewBuffer([]byte(`{}`)), isErr: true, }, { name: "empty documentPaths array", api: "http://127.0.0.1:5000/api/tool/firestore-get-docs/invoke", requestBody: bytes.NewBuffer([]byte(`{"documentPaths": []}`)), isErr: true, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %v", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if tc.wantRegex != "" { matched, err := regexp.MatchString(tc.wantRegex, got) if err != nil { t.Fatalf("invalid regex pattern: %v", err) } if !matched { t.Fatalf("result does not match expected pattern.\nGot: %s\nWant pattern: %s", got, tc.wantRegex) } } }) } } func runFirestoreListCollectionsTest(t *testing.T, collectionName, subCollectionName, parentDocPath string) { invokeTcs := []struct { name string api string requestBody io.Reader want string isErr bool }{ { name: "list root collections", api: "http://127.0.0.1:5000/api/tool/firestore-list-colls/invoke", requestBody: bytes.NewBuffer([]byte(`{}`)), want: collectionName, isErr: false, }, { name: "list subcollections", api: "http://127.0.0.1:5000/api/tool/firestore-list-colls/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"parentPath": "%s"}`, parentDocPath))), want: subCollectionName, isErr: false, }, { name: "list collections for non-existent parent", api: "http://127.0.0.1:5000/api/tool/firestore-list-colls/invoke", requestBody: bytes.NewBuffer([]byte(`{"parentPath": "non-existent-collection/non-existent-doc"}`)), want: `[]`, // Empty array for no collections isErr: false, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %v", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if !strings.Contains(got, tc.want) { t.Fatalf("expected %q to contain %q, but it did not", got, tc.want) } }) } } func runFirestoreDeleteDocumentsTest(t *testing.T, docPath string) { invokeTcs := []struct { name string api string requestBody io.Reader want string isErr bool }{ { name: "delete single document", api: "http://127.0.0.1:5000/api/tool/firestore-delete-docs/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"documentPaths": ["%s"]}`, docPath))), want: `"success":true`, isErr: false, }, { name: "delete non-existent document", api: "http://127.0.0.1:5000/api/tool/firestore-delete-docs/invoke", requestBody: bytes.NewBuffer([]byte(`{"documentPaths": ["non-existent-collection/non-existent-doc"]}`)), want: `"success":true`, // Firestore delete succeeds even if doc doesn't exist isErr: false, }, { name: "missing documentPaths parameter", api: "http://127.0.0.1:5000/api/tool/firestore-delete-docs/invoke", requestBody: bytes.NewBuffer([]byte(`{}`)), isErr: true, }, { name: "empty documentPaths array", api: "http://127.0.0.1:5000/api/tool/firestore-delete-docs/invoke", requestBody: bytes.NewBuffer([]byte(`{"documentPaths": []}`)), isErr: true, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %v", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if !strings.Contains(got, tc.want) { t.Fatalf("expected %q to contain %q, but it did not", got, tc.want) } }) } } func runFirestoreQueryTest(t *testing.T, collectionName string) { invokeTcs := []struct { name string api string requestBody io.Reader wantRegex string isErr bool }{ { name: "query with parameterized filters - age greater than", api: "http://127.0.0.1:5000/api/tool/firestore-query-param/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collection": "%s", "operator": ">", "ageValue": "25" }`, collectionName))), wantRegex: `"name":"Alice"`, isErr: false, }, { name: "query with parameterized filters - exact name match", api: "http://127.0.0.1:5000/api/tool/firestore-query-param/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collection": "%s", "operator": "==", "ageValue": "25" }`, collectionName))), wantRegex: `"name":"Bob"`, isErr: false, }, { name: "query with parameterized filters - age less than or equal", api: "http://127.0.0.1:5000/api/tool/firestore-query-param/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collection": "%s", "operator": "<=", "ageValue": "29" }`, collectionName))), wantRegex: `"name":"Bob"`, isErr: false, }, { name: "missing required parameter", api: "http://127.0.0.1:5000/api/tool/firestore-query-param/invoke", requestBody: bytes.NewBuffer([]byte(`{"collection": "test", "operator": ">"}`)), isErr: true, }, { name: "query non-existent collection with parameters", api: "http://127.0.0.1:5000/api/tool/firestore-query-param/invoke", requestBody: bytes.NewBuffer([]byte(`{ "collection": "non-existent-collection", "operator": "==", "ageValue": "30" }`)), wantRegex: `^\[\]$`, // Empty array isErr: false, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %v", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if tc.wantRegex != "" { matched, err := regexp.MatchString(tc.wantRegex, got) if err != nil { t.Fatalf("invalid regex pattern: %v", err) } if !matched { t.Fatalf("result does not match expected pattern.\nGot: %s\nWant pattern: %s", got, tc.wantRegex) } } }) } } func runFirestoreQuerySelectArrayTest(t *testing.T, collectionName string) { invokeTcs := []struct { name string api string requestBody io.Reader wantRegex string validateFields bool isErr bool }{ { name: "query with array select fields - single field", api: "http://127.0.0.1:5000/api/tool/firestore-query-select-array/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collection": "%s", "fields": ["name"] }`, collectionName))), wantRegex: `"name":"`, validateFields: true, isErr: false, }, { name: "query with array select fields - multiple fields", api: "http://127.0.0.1:5000/api/tool/firestore-query-select-array/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collection": "%s", "fields": ["name", "age"] }`, collectionName))), wantRegex: `"name":".*"age":`, validateFields: true, isErr: false, }, { name: "query with empty array select fields", api: "http://127.0.0.1:5000/api/tool/firestore-query-select-array/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collection": "%s", "fields": [] }`, collectionName))), wantRegex: `\[.*\]`, // Should return documents with all fields isErr: false, }, { name: "missing fields parameter", api: "http://127.0.0.1:5000/api/tool/firestore-query-select-array/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"collection": "%s"}`, collectionName))), isErr: true, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %v", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if tc.wantRegex != "" { matched, err := regexp.MatchString(tc.wantRegex, got) if err != nil { t.Fatalf("invalid regex pattern: %v", err) } if !matched { t.Fatalf("result does not match expected pattern.\nGot: %s\nWant pattern: %s", got, tc.wantRegex) } } // Additional validation for field selection if tc.validateFields { // Parse the result to check if only selected fields are present var results []map[string]interface{} err = json.Unmarshal([]byte(got), &results) if err != nil { t.Fatalf("error parsing result as JSON array: %v", err) } // For single field test, ensure only 'name' field is present in data if tc.name == "query with array select fields - single field" && len(results) > 0 { for _, result := range results { if data, ok := result["data"].(map[string]interface{}); ok { if _, hasName := data["name"]; !hasName { t.Fatalf("expected 'name' field in data, but not found") } // The 'age' field should not be present when only 'name' is selected if _, hasAge := data["age"]; hasAge { t.Fatalf("unexpected 'age' field in data when only 'name' was selected") } } } } // For multiple fields test, ensure both fields are present if tc.name == "query with array select fields - multiple fields" && len(results) > 0 { for _, result := range results { if data, ok := result["data"].(map[string]interface{}); ok { if _, hasName := data["name"]; !hasName { t.Fatalf("expected 'name' field in data, but not found") } if _, hasAge := data["age"]; !hasAge { t.Fatalf("expected 'age' field in data, but not found") } } } } } }) } } func runFirestoreQueryCollectionTest(t *testing.T, collectionName string) { invokeTcs := []struct { name string api string requestBody io.Reader wantRegex string isErr bool }{ { name: "query collection with filter", api: "http://127.0.0.1:5000/api/tool/firestore-query-coll/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collectionPath": "%s", "filters": ["{\"field\": \"age\", \"op\": \">\", \"value\": 25}"], "orderBy": "", "limit": 10 }`, collectionName))), wantRegex: `"name":"Alice"`, isErr: false, }, { name: "query collection with orderBy", api: "http://127.0.0.1:5000/api/tool/firestore-query-coll/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collectionPath": "%s", "filters": [], "orderBy": "{\"field\": \"age\", \"direction\": \"DESCENDING\"}", "limit": 2 }`, collectionName))), wantRegex: `"age":35.*"age":30`, // Should be ordered by age descending (Charlie=35, Alice=30) isErr: false, }, { name: "query collection with multiple filters", api: "http://127.0.0.1:5000/api/tool/firestore-query-coll/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collectionPath": "%s", "filters": [ "{\"field\": \"age\", \"op\": \">=\", \"value\": 25}", "{\"field\": \"age\", \"op\": \"<=\", \"value\": 30}" ], "orderBy": "", "limit": 10 }`, collectionName))), wantRegex: `"name":"Bob".*"name":"Alice"`, // Results may be ordered by document ID isErr: false, }, { name: "query with limit", api: "http://127.0.0.1:5000/api/tool/firestore-query-coll/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collectionPath": "%s", "filters": [], "orderBy": "", "limit": 1 }`, collectionName))), wantRegex: `^\[{.*}\]$`, // Should return exactly one document isErr: false, }, { name: "query non-existent collection", api: "http://127.0.0.1:5000/api/tool/firestore-query-coll/invoke", requestBody: bytes.NewBuffer([]byte(`{ "collectionPath": "non-existent-collection", "filters": [], "orderBy": "", "limit": 10 }`)), wantRegex: `^\[\]$`, // Empty array isErr: false, }, { name: "missing collectionPath parameter", api: "http://127.0.0.1:5000/api/tool/firestore-query-coll/invoke", requestBody: bytes.NewBuffer([]byte(`{}`)), isErr: true, }, { name: "invalid filter operator", api: "http://127.0.0.1:5000/api/tool/firestore-query-coll/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collectionPath": "%s", "filters": ["{\"field\": \"age\", \"op\": \"INVALID\", \"value\": 25}"], "orderBy": "" }`, collectionName))), isErr: true, }, { name: "query with analyzeQuery", api: "http://127.0.0.1:5000/api/tool/firestore-query-coll/invoke", requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{ "collectionPath": "%s", "filters": [], "orderBy": "", "analyzeQuery": true, "limit": 1 }`, collectionName))), wantRegex: `"documents":\[.*\]`, isErr: false, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %v", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if tc.wantRegex != "" { matched, err := regexp.MatchString(tc.wantRegex, got) if err != nil { t.Fatalf("invalid regex pattern: %v", err) } if !matched { t.Fatalf("result does not match expected pattern.\nGot: %s\nWant pattern: %s", got, tc.wantRegex) } } }) } } ``` -------------------------------------------------------------------------------- /tests/tool.go: -------------------------------------------------------------------------------- ```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 tests import ( "bytes" "context" "database/sql" "encoding/json" "fmt" "io" "net/http" "reflect" "sort" "strings" "sync" "testing" "time" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" "github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc" "github.com/googleapis/genai-toolbox/internal/sources" ) // RunToolGet runs the tool get endpoint func RunToolGetTest(t *testing.T) { // Test tool get endpoint tcs := []struct { name string api string want map[string]any }{ { name: "get my-simple-tool", api: "http://127.0.0.1:5000/api/tool/my-simple-tool/", want: map[string]any{ "my-simple-tool": map[string]any{ "description": "Simple tool to test end to end functionality.", "parameters": []any{}, "authRequired": []any{}, }, }, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { resp, err := http.Get(tc.api) if err != nil { t.Fatalf("error when sending a request: %s", err) } defer resp.Body.Close() if resp.StatusCode != 200 { t.Fatalf("response status code is not 200") } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body") } got, ok := body["tools"] if !ok { t.Fatalf("unable to find tools in response body") } if !reflect.DeepEqual(got, tc.want) { t.Fatalf("got %q, want %q", got, tc.want) } }) } } func RunToolGetTestByName(t *testing.T, name string, want map[string]any) { // Test tool get endpoint tcs := []struct { name string api string want map[string]any }{ { name: fmt.Sprintf("get %s", name), api: fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/", name), want: want, }, } for _, tc := range tcs { t.Run(tc.name, func(t *testing.T) { resp, err := http.Get(tc.api) if err != nil { t.Fatalf("error when sending a request: %s", err) } defer resp.Body.Close() if resp.StatusCode != 200 { t.Fatalf("response status code is not 200") } var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body") } got, ok := body["tools"] if !ok { t.Fatalf("unable to find tools in response body") } if !reflect.DeepEqual(got, tc.want) { t.Fatalf("got %q, want %q", got, tc.want) } }) } } // RunToolInvokeSimpleTest runs the tool invoke endpoint with no parameters func RunToolInvokeSimpleTest(t *testing.T, name string, simpleWant string) { // Test tool invoke endpoint invokeTcs := []struct { name string api string requestHeader map[string]string requestBody io.Reader want string isErr bool }{ { name: fmt.Sprintf("invoke %s", name), api: fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/invoke", name), requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{}`)), want: simpleWant, isErr: false, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { // Send Tool invocation request req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") for k, v := range tc.requestHeader { req.Header.Add(k, v) } resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } // Check response body var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body") } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if !strings.Contains(got, tc.want) { t.Fatalf("unexpected value: got %q, want %q", got, tc.want) } }) } } func RunToolInvokeParametersTest(t *testing.T, name string, params []byte, simpleWant string) { // Test tool invoke endpoint invokeTcs := []struct { name string api string requestHeader map[string]string requestBody io.Reader want string isErr bool }{ { name: fmt.Sprintf("invoke %s", name), api: fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/invoke", name), requestHeader: map[string]string{}, requestBody: bytes.NewBuffer(params), want: simpleWant, isErr: false, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { // Send Tool invocation request req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") for k, v := range tc.requestHeader { req.Header.Add(k, v) } resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } // Check response body var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body") } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if !strings.Contains(got, tc.want) { t.Fatalf("unexpected value: got %q, want %q", got, tc.want) } }) } } // RunToolInvoke runs the tool invoke endpoint func RunToolInvokeTest(t *testing.T, select1Want string, options ...InvokeTestOption) { // Resolve options // Default values for InvokeTestConfig configs := &InvokeTestConfig{ myToolId3NameAliceWant: "[{\"id\":1,\"name\":\"Alice\"},{\"id\":3,\"name\":\"Sid\"}]", myToolById4Want: "[{\"id\":4,\"name\":null}]", myArrayToolWant: "[{\"id\":1,\"name\":\"Alice\"},{\"id\":3,\"name\":\"Sid\"}]", nullWant: "null", supportOptionalNullParam: true, supportArrayParam: true, supportClientAuth: false, supportSelect1Want: true, supportSelect1Auth: true, } // Apply provided options for _, option := range options { option(configs) } // Get ID token idToken, err := GetGoogleIdToken(ClientId) if err != nil { t.Fatalf("error getting Google ID token: %s", err) } // Get access token accessToken, err := sources.GetIAMAccessToken(t.Context()) if err != nil { t.Fatalf("error getting access token from ADC: %s", err) } accessToken = "Bearer " + accessToken // Test tool invoke endpoint invokeTcs := []struct { name string api string enabled bool requestHeader map[string]string requestBody io.Reader wantStatusCode int wantBody string }{ { name: "invoke my-simple-tool", api: "http://127.0.0.1:5000/api/tool/my-simple-tool/invoke", enabled: configs.supportSelect1Want, requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantBody: select1Want, wantStatusCode: http.StatusOK, }, { name: "invoke my-tool", api: "http://127.0.0.1:5000/api/tool/my-tool/invoke", enabled: true, requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{"id": 3, "name": "Alice"}`)), wantBody: configs.myToolId3NameAliceWant, wantStatusCode: http.StatusOK, }, { name: "invoke my-tool-by-id with nil response", api: "http://127.0.0.1:5000/api/tool/my-tool-by-id/invoke", enabled: true, requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{"id": 4}`)), wantBody: configs.myToolById4Want, wantStatusCode: http.StatusOK, }, { name: "invoke my-tool-by-name with nil response", api: "http://127.0.0.1:5000/api/tool/my-tool-by-name/invoke", enabled: configs.supportOptionalNullParam, requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantBody: configs.nullWant, wantStatusCode: http.StatusOK, }, { name: "Invoke my-tool without parameters", api: "http://127.0.0.1:5000/api/tool/my-tool/invoke", enabled: true, requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantBody: "", wantStatusCode: http.StatusBadRequest, }, { name: "Invoke my-tool with insufficient parameters", api: "http://127.0.0.1:5000/api/tool/my-tool/invoke", enabled: true, requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{"id": 1}`)), wantBody: "", wantStatusCode: http.StatusBadRequest, }, { name: "invoke my-array-tool", api: "http://127.0.0.1:5000/api/tool/my-array-tool/invoke", enabled: configs.supportArrayParam, requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{"idArray": [1,2,3], "nameArray": ["Alice", "Sid", "RandomName"], "cmdArray": ["HGETALL", "row3"]}`)), wantBody: configs.myArrayToolWant, wantStatusCode: http.StatusOK, }, { name: "Invoke my-auth-tool with auth token", api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke", enabled: configs.supportSelect1Auth, requestHeader: map[string]string{"my-google-auth_token": idToken}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantBody: "[{\"name\":\"Alice\"}]", wantStatusCode: http.StatusOK, }, { name: "Invoke my-auth-tool with invalid auth token", api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke", enabled: configs.supportSelect1Auth, requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantStatusCode: http.StatusUnauthorized, }, { name: "Invoke my-auth-tool without auth token", api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke", enabled: true, requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantStatusCode: http.StatusUnauthorized, }, { name: "Invoke my-auth-required-tool with auth token", api: "http://127.0.0.1:5000/api/tool/my-auth-required-tool/invoke", enabled: configs.supportSelect1Auth, requestHeader: map[string]string{"my-google-auth_token": idToken}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantBody: select1Want, wantStatusCode: http.StatusOK, }, { name: "Invoke my-auth-required-tool with invalid auth token", api: "http://127.0.0.1:5000/api/tool/my-auth-required-tool/invoke", enabled: true, requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantStatusCode: http.StatusUnauthorized, }, { name: "Invoke my-auth-required-tool without auth token", api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke", enabled: true, requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantStatusCode: http.StatusUnauthorized, }, { name: "Invoke my-client-auth-tool with auth token", api: "http://127.0.0.1:5000/api/tool/my-client-auth-tool/invoke", enabled: configs.supportClientAuth, requestHeader: map[string]string{"Authorization": accessToken}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantBody: select1Want, wantStatusCode: http.StatusOK, }, { name: "Invoke my-client-auth-tool without auth token", api: "http://127.0.0.1:5000/api/tool/my-client-auth-tool/invoke", enabled: configs.supportClientAuth, requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantStatusCode: http.StatusUnauthorized, }, { name: "Invoke my-client-auth-tool with invalid auth token", api: "http://127.0.0.1:5000/api/tool/my-client-auth-tool/invoke", enabled: configs.supportClientAuth, requestHeader: map[string]string{"Authorization": "Bearer invalid-token"}, requestBody: bytes.NewBuffer([]byte(`{}`)), wantStatusCode: http.StatusUnauthorized, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { if !tc.enabled { return } // Send Tool invocation request req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") // Add headers for k, v := range tc.requestHeader { req.Header.Add(k, v) } resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() // Check status code if resp.StatusCode != tc.wantStatusCode { body, _ := io.ReadAll(resp.Body) t.Errorf("StatusCode mismatch: got %d, want %d. Response body: %s", resp.StatusCode, tc.wantStatusCode, string(body)) } // skip response body check if tc.wantBody == "" { return } // Check response body var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body: %s", err) } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if got != tc.wantBody { t.Fatalf("unexpected value: got %q, want %q", got, tc.wantBody) } }) } } // RunToolInvokeWithTemplateParameters runs tool invoke test cases with template parameters. func RunToolInvokeWithTemplateParameters(t *testing.T, tableName string, options ...TemplateParamOption) { // Resolve options // Default values for TemplateParameterTestConfig configs := &TemplateParameterTestConfig{ ddlWant: "null", selectAllWant: "[{\"age\":21,\"id\":1,\"name\":\"Alex\"},{\"age\":100,\"id\":2,\"name\":\"Alice\"}]", selectId1Want: "[{\"age\":21,\"id\":1,\"name\":\"Alex\"}]", selectNameWant: "[{\"age\":21,\"id\":1,\"name\":\"Alex\"}]", selectEmptyWant: "null", insert1Want: "null", nameFieldArray: `["name"]`, nameColFilter: "name", createColArray: `["id INT","name VARCHAR(20)","age INT"]`, supportDdl: true, supportInsert: true, } // Apply provided options for _, option := range options { option(configs) } selectOnlyNamesWant := "[{\"name\":\"Alex\"},{\"name\":\"Alice\"}]" // Test tool invoke endpoint invokeTcs := []struct { name string enabled bool ddl bool insert bool api string requestHeader map[string]string requestBody io.Reader want string isErr bool }{ { name: "invoke create-table-templateParams-tool", ddl: true, api: "http://127.0.0.1:5000/api/tool/create-table-templateParams-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s", "columns":%s}`, tableName, configs.createColArray))), want: configs.ddlWant, isErr: false, }, { name: "invoke insert-table-templateParams-tool", insert: true, api: "http://127.0.0.1:5000/api/tool/insert-table-templateParams-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s", "columns":["id","name","age"], "values":"1, 'Alex', 21"}`, tableName))), want: configs.insert1Want, isErr: false, }, { name: "invoke insert-table-templateParams-tool", insert: true, api: "http://127.0.0.1:5000/api/tool/insert-table-templateParams-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s", "columns":["id","name","age"], "values":"2, 'Alice', 100"}`, tableName))), want: configs.insert1Want, isErr: false, }, { name: "invoke select-templateParams-tool", api: "http://127.0.0.1:5000/api/tool/select-templateParams-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s"}`, tableName))), want: configs.selectAllWant, isErr: false, }, { name: "invoke select-templateParams-combined-tool", api: "http://127.0.0.1:5000/api/tool/select-templateParams-combined-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"id": 1, "tableName": "%s"}`, tableName))), want: configs.selectId1Want, isErr: false, }, { name: "invoke select-templateParams-combined-tool with no results", api: "http://127.0.0.1:5000/api/tool/select-templateParams-combined-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"id": 999, "tableName": "%s"}`, tableName))), want: configs.selectEmptyWant, isErr: false, }, { name: "invoke select-fields-templateParams-tool", enabled: configs.supportSelectFields, api: "http://127.0.0.1:5000/api/tool/select-fields-templateParams-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s", "fields":%s}`, tableName, configs.nameFieldArray))), want: selectOnlyNamesWant, isErr: false, }, { name: "invoke select-filter-templateParams-combined-tool", api: "http://127.0.0.1:5000/api/tool/select-filter-templateParams-combined-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"name": "Alex", "tableName": "%s", "columnFilter": "%s"}`, tableName, configs.nameColFilter))), want: configs.selectNameWant, isErr: false, }, { name: "invoke drop-table-templateParams-tool", ddl: true, api: "http://127.0.0.1:5000/api/tool/drop-table-templateParams-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s"}`, tableName))), want: configs.ddlWant, isErr: false, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { if !tc.enabled { return } // if test case is DDL and source support ddl test cases ddlAllow := !tc.ddl || (tc.ddl && configs.supportDdl) // if test case is insert statement and source support insert test cases insertAllow := !tc.insert || (tc.insert && configs.supportInsert) if ddlAllow && insertAllow { // Send Tool invocation request req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") for k, v := range tc.requestHeader { req.Header.Add(k, v) } resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } // Check response body var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body") } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if got != tc.want { t.Fatalf("unexpected value: got %q, want %q", got, tc.want) } } }) } } func RunExecuteSqlToolInvokeTest(t *testing.T, createTableStatement, select1Want string, options ...ExecuteSqlOption) { // Resolve options // Default values for ExecuteSqlTestConfig configs := &ExecuteSqlTestConfig{ select1Statement: `"SELECT 1"`, } // Apply provided options for _, option := range options { option(configs) } // Get ID token idToken, err := GetGoogleIdToken(ClientId) if err != nil { t.Fatalf("error getting Google ID token: %s", err) } // Test tool invoke endpoint invokeTcs := []struct { name string api string requestHeader map[string]string requestBody io.Reader want string isErr bool }{ { name: "invoke my-exec-sql-tool", api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sql": %s}`, configs.select1Statement))), want: select1Want, isErr: false, }, { name: "invoke my-exec-sql-tool create table", api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sql": %s}`, createTableStatement))), want: "null", isErr: false, }, { name: "invoke my-exec-sql-tool select table", api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT * FROM t"}`)), want: "null", isErr: false, }, { name: "invoke my-exec-sql-tool drop table", api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{"sql":"DROP TABLE t"}`)), want: "null", isErr: false, }, { name: "invoke my-exec-sql-tool without body", api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(`{}`)), isErr: true, }, { name: "Invoke my-auth-exec-sql-tool with auth token", api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke", requestHeader: map[string]string{"my-google-auth_token": idToken}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sql": %s}`, configs.select1Statement))), isErr: false, want: select1Want, }, { name: "Invoke my-auth-exec-sql-tool with invalid auth token", api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke", requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sql": %s}`, configs.select1Statement))), isErr: true, }, { name: "Invoke my-auth-exec-sql-tool without auth token", api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke", requestHeader: map[string]string{}, requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sql": %s}`, configs.select1Statement))), isErr: true, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { // Send Tool invocation request req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Add("Content-type", "application/json") for k, v := range tc.requestHeader { req.Header.Add(k, v) } resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { if tc.isErr { return } bodyBytes, _ := io.ReadAll(resp.Body) t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes)) } // Check response body var body map[string]interface{} err = json.NewDecoder(resp.Body).Decode(&body) if err != nil { t.Fatalf("error parsing response body") } got, ok := body["result"].(string) if !ok { t.Fatalf("unable to find result in response body") } if got != tc.want { t.Fatalf("unexpected value: got %q, want %q", got, tc.want) } }) } } // RunInitialize runs the initialize lifecycle for mcp to set up client-server connection func RunInitialize(t *testing.T, protocolVersion string) string { url := "http://127.0.0.1:5000/mcp" initializeRequestBody := map[string]any{ "jsonrpc": "2.0", "id": "mcp-initialize", "method": "initialize", "params": map[string]any{ "protocolVersion": protocolVersion, }, } reqMarshal, err := json.Marshal(initializeRequestBody) if err != nil { t.Fatalf("unexpected error during marshaling of body") } resp, _ := RunRequest(t, http.MethodPost, url, bytes.NewBuffer(reqMarshal), nil) if resp.StatusCode != 200 { t.Fatalf("response status code is not 200") } if contentType := resp.Header.Get("Content-type"); contentType != "application/json" { t.Fatalf("unexpected content-type header: want %s, got %s", "application/json", contentType) } sessionId := resp.Header.Get("Mcp-Session-Id") header := map[string]string{} if sessionId != "" { header["Mcp-Session-Id"] = sessionId } initializeNotificationBody := map[string]any{ "jsonrpc": "2.0", "method": "notifications/initialized", } notiMarshal, err := json.Marshal(initializeNotificationBody) if err != nil { t.Fatalf("unexpected error during marshaling of notifications body") } _, _ = RunRequest(t, http.MethodPost, url, bytes.NewBuffer(notiMarshal), header) return sessionId } // RunMCPToolCallMethod runs the tool/call for mcp endpoint func RunMCPToolCallMethod(t *testing.T, myFailToolWant, select1Want string, options ...McpTestOption) { // Resolve options // Default values for MCPTestConfig configs := &MCPTestConfig{ myToolId3NameAliceWant: `{"jsonrpc":"2.0","id":"my-tool","result":{"content":[{"type":"text","text":"{\"id\":1,\"name\":\"Alice\"}"},{"type":"text","text":"{\"id\":3,\"name\":\"Sid\"}"}]}}`, supportClientAuth: false, supportSelect1Auth: true, } // Apply provided options for _, option := range options { option(configs) } sessionId := RunInitialize(t, "2024-11-05") // Get access token accessToken, err := sources.GetIAMAccessToken(t.Context()) if err != nil { t.Fatalf("error getting access token from ADC: %s", err) } accessToken = "Bearer " + accessToken idToken, err := GetGoogleIdToken(ClientId) if err != nil { t.Fatalf("error getting Google ID token: %s", err) } // Test tool invoke endpoint invokeTcs := []struct { name string api string enabled bool // switch to turn on/off the test case requestBody jsonrpc.JSONRPCRequest requestHeader map[string]string wantStatusCode int wantBody string }{ { name: "MCP Invoke my-tool", api: "http://127.0.0.1:5000/mcp", enabled: true, requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "my-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-tool", "arguments": map[string]any{ "id": int(3), "name": "Alice", }, }, }, wantStatusCode: http.StatusOK, wantBody: configs.myToolId3NameAliceWant, }, { name: "MCP Invoke invalid tool", api: "http://127.0.0.1:5000/mcp", enabled: true, requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invalid-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "foo", "arguments": map[string]any{}, }, }, wantStatusCode: http.StatusOK, wantBody: `{"jsonrpc":"2.0","id":"invalid-tool","error":{"code":-32602,"message":"invalid tool name: tool with name \"foo\" does not exist"}}`, }, { name: "MCP Invoke my-tool without parameters", api: "http://127.0.0.1:5000/mcp", enabled: true, requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke-without-parameter", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-tool", "arguments": map[string]any{}, }, }, wantStatusCode: http.StatusOK, wantBody: `{"jsonrpc":"2.0","id":"invoke-without-parameter","error":{"code":-32602,"message":"provided parameters were invalid: parameter \"id\" is required"}}`, }, { name: "MCP Invoke my-tool with insufficient parameters", api: "http://127.0.0.1:5000/mcp", enabled: true, requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke-insufficient-parameter", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-tool", "arguments": map[string]any{"id": 1}, }, }, wantStatusCode: http.StatusOK, wantBody: `{"jsonrpc":"2.0","id":"invoke-insufficient-parameter","error":{"code":-32602,"message":"provided parameters were invalid: parameter \"name\" is required"}}`, }, { name: "MCP Invoke my-auth-required-tool", api: "http://127.0.0.1:5000/mcp", enabled: configs.supportSelect1Auth, requestHeader: map[string]string{"my-google-auth_token": idToken}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke my-auth-required-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-auth-required-tool", "arguments": map[string]any{}, }, }, wantStatusCode: http.StatusOK, wantBody: select1Want, }, { name: "MCP Invoke my-auth-required-tool with invalid auth token", api: "http://127.0.0.1:5000/mcp", requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke my-auth-required-tool with invalid token", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-auth-required-tool", "arguments": map[string]any{}, }, }, wantStatusCode: http.StatusUnauthorized, wantBody: "{\"jsonrpc\":\"2.0\",\"id\":\"invoke my-auth-required-tool with invalid token\",\"error\":{\"code\":-32600,\"message\":\"unauthorized Tool call: Please make sure your specify correct auth headers: unauthorized\"}}", }, { name: "MCP Invoke my-auth-required-tool without auth token", api: "http://127.0.0.1:5000/mcp", requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke my-auth-required-tool without token", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-auth-required-tool", "arguments": map[string]any{}, }, }, wantStatusCode: http.StatusUnauthorized, wantBody: "{\"jsonrpc\":\"2.0\",\"id\":\"invoke my-auth-required-tool without token\",\"error\":{\"code\":-32600,\"message\":\"unauthorized Tool call: Please make sure your specify correct auth headers: unauthorized\"}}", }, { name: "MCP Invoke my-client-auth-tool", enabled: configs.supportClientAuth, api: "http://127.0.0.1:5000/mcp", requestHeader: map[string]string{"Authorization": accessToken}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke my-client-auth-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-client-auth-tool", "arguments": map[string]any{}, }, }, wantStatusCode: http.StatusOK, wantBody: "{\"jsonrpc\":\"2.0\",\"id\":\"invoke my-client-auth-tool\",\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"{\\\"f0_\\\":1}\"}]}}", }, { name: "MCP Invoke my-client-auth-tool without access token", enabled: configs.supportClientAuth, api: "http://127.0.0.1:5000/mcp", requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke my-client-auth-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-client-auth-tool", "arguments": map[string]any{}, }, }, wantStatusCode: http.StatusUnauthorized, wantBody: "{\"jsonrpc\":\"2.0\",\"id\":\"invoke my-client-auth-tool\",\"error\":{\"code\":-32600,\"message\":\"missing access token in the 'Authorization' header\"}", }, { name: "MCP Invoke my-client-auth-tool with invalid access token", enabled: configs.supportClientAuth, api: "http://127.0.0.1:5000/mcp", requestHeader: map[string]string{"Authorization": "Bearer invalid-token"}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke my-client-auth-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-client-auth-tool", "arguments": map[string]any{}, }, }, wantStatusCode: http.StatusUnauthorized, }, { name: "MCP Invoke my-fail-tool", api: "http://127.0.0.1:5000/mcp", enabled: true, requestHeader: map[string]string{}, requestBody: jsonrpc.JSONRPCRequest{ Jsonrpc: "2.0", Id: "invoke-fail-tool", Request: jsonrpc.Request{ Method: "tools/call", }, Params: map[string]any{ "name": "my-fail-tool", "arguments": map[string]any{"id": 1}, }, }, wantStatusCode: http.StatusOK, wantBody: myFailToolWant, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { if !tc.enabled { return } reqMarshal, err := json.Marshal(tc.requestBody) if err != nil { t.Fatalf("unexpected error during marshaling of request body") } // add headers headers := map[string]string{} if sessionId != "" { headers["Mcp-Session-Id"] = sessionId } for key, value := range tc.requestHeader { headers[key] = value } httpResponse, respBody := RunRequest(t, http.MethodPost, tc.api, bytes.NewBuffer(reqMarshal), headers) // Check status code if httpResponse.StatusCode != tc.wantStatusCode { t.Errorf("StatusCode mismatch: got %d, want %d", httpResponse.StatusCode, tc.wantStatusCode) } // Check response body got := string(bytes.TrimSpace(respBody)) if !strings.Contains(got, tc.wantBody) { t.Fatalf("Expected substring not found:\ngot: %q\nwant: %q (to be contained within got)", got, tc.wantBody) } }) } } // RunMySQLListTablesTest run tests against the mysql-list-tables tool func RunMySQLListTablesTest(t *testing.T, databaseName, tableNameParam, tableNameAuth string) { type tableInfo struct { ObjectName string `json:"object_name"` SchemaName string `json:"schema_name"` ObjectDetails string `json:"object_details"` } type column struct { DataType string `json:"data_type"` ColumnName string `json:"column_name"` ColumnComment string `json:"column_comment"` ColumnDefault any `json:"column_default"` IsNotNullable int `json:"is_not_nullable"` OrdinalPosition int `json:"ordinal_position"` } type objectDetails struct { Owner any `json:"owner"` Columns []column `json:"columns"` Comment string `json:"comment"` Indexes []any `json:"indexes"` Triggers []any `json:"triggers"` Constraints []any `json:"constraints"` ObjectName string `json:"object_name"` ObjectType string `json:"object_type"` SchemaName string `json:"schema_name"` } paramTableWant := objectDetails{ ObjectName: tableNameParam, SchemaName: databaseName, ObjectType: "TABLE", Columns: []column{ {DataType: "int", ColumnName: "id", IsNotNullable: 1, OrdinalPosition: 1}, {DataType: "varchar(255)", ColumnName: "name", OrdinalPosition: 2}, }, Indexes: []any{map[string]any{"index_columns": []any{"id"}, "index_name": "PRIMARY", "is_primary": float64(1), "is_unique": float64(1)}}, Triggers: []any{}, Constraints: []any{map[string]any{"constraint_columns": []any{"id"}, "constraint_name": "PRIMARY", "constraint_type": "PRIMARY KEY", "foreign_key_referenced_columns": any(nil), "foreign_key_referenced_table": any(nil), "constraint_definition": ""}}, } authTableWant := objectDetails{ ObjectName: tableNameAuth, SchemaName: databaseName, ObjectType: "TABLE", Columns: []column{ {DataType: "int", ColumnName: "id", IsNotNullable: 1, OrdinalPosition: 1}, {DataType: "varchar(255)", ColumnName: "name", OrdinalPosition: 2}, {DataType: "varchar(255)", ColumnName: "email", OrdinalPosition: 3}, }, Indexes: []any{map[string]any{"index_columns": []any{"id"}, "index_name": "PRIMARY", "is_primary": float64(1), "is_unique": float64(1)}}, Triggers: []any{}, Constraints: []any{map[string]any{"constraint_columns": []any{"id"}, "constraint_name": "PRIMARY", "constraint_type": "PRIMARY KEY", "foreign_key_referenced_columns": any(nil), "foreign_key_referenced_table": any(nil), "constraint_definition": ""}}, } invokeTcs := []struct { name string requestBody io.Reader wantStatusCode int want any isSimple bool isAllTables bool }{ { name: "invoke list_tables for all tables detailed output", requestBody: bytes.NewBufferString(`{"table_names":""}`), wantStatusCode: http.StatusOK, want: []objectDetails{authTableWant, paramTableWant}, isAllTables: true, }, { name: "invoke list_tables detailed output", requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_names": "%s"}`, tableNameAuth)), wantStatusCode: http.StatusOK, want: []objectDetails{authTableWant}, }, { name: "invoke list_tables simple output", requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_names": "%s", "output_format": "simple"}`, tableNameAuth)), wantStatusCode: http.StatusOK, want: []map[string]any{{"name": tableNameAuth}}, isSimple: true, }, { name: "invoke list_tables with multiple table names", requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_names": "%s,%s"}`, tableNameParam, tableNameAuth)), wantStatusCode: http.StatusOK, want: []objectDetails{authTableWant, paramTableWant}, }, { name: "invoke list_tables with one existing and one non-existent table", requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_names": "%s,non_existent_table"}`, tableNameAuth)), wantStatusCode: http.StatusOK, want: []objectDetails{authTableWant}, }, { name: "invoke list_tables with non-existent table", requestBody: bytes.NewBufferString(`{"table_names": "non_existent_table"}`), wantStatusCode: http.StatusOK, want: nil, }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { const api = "http://127.0.0.1:5000/api/tool/list_tables/invoke" req, err := http.NewRequest(http.MethodPost, api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %v", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %v", err) } defer resp.Body.Close() if resp.StatusCode != tc.wantStatusCode { body, _ := io.ReadAll(resp.Body) t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(body)) } if tc.wantStatusCode != http.StatusOK { return } var bodyWrapper struct { Result json.RawMessage `json:"result"` } if err := json.NewDecoder(resp.Body).Decode(&bodyWrapper); err != nil { t.Fatalf("error decoding response wrapper: %v", err) } var resultString string if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil { resultString = string(bodyWrapper.Result) } var got any if tc.isSimple { var tables []tableInfo if err := json.Unmarshal([]byte(resultString), &tables); err != nil { t.Fatalf("failed to unmarshal outer JSON array into []tableInfo: %v", err) } var details []map[string]any for _, table := range tables { var d map[string]any if err := json.Unmarshal([]byte(table.ObjectDetails), &d); err != nil { t.Fatalf("failed to unmarshal nested ObjectDetails string: %v", err) } details = append(details, d) } got = details } else { if resultString == "null" { got = nil } else { var tables []tableInfo if err := json.Unmarshal([]byte(resultString), &tables); err != nil { t.Fatalf("failed to unmarshal outer JSON array into []tableInfo: %v", err) } var details []objectDetails for _, table := range tables { var d objectDetails if err := json.Unmarshal([]byte(table.ObjectDetails), &d); err != nil { t.Fatalf("failed to unmarshal nested ObjectDetails string: %v", err) } details = append(details, d) } got = details } } opts := []cmp.Option{ cmpopts.SortSlices(func(a, b objectDetails) bool { return a.ObjectName < b.ObjectName }), cmpopts.SortSlices(func(a, b column) bool { return a.ColumnName < b.ColumnName }), cmpopts.SortSlices(func(a, b map[string]any) bool { return a["name"].(string) < b["name"].(string) }), } // Checking only the current database where the test tables are created to avoid brittle tests. if tc.isAllTables { var filteredGot []objectDetails if got != nil { for _, item := range got.([]objectDetails) { if item.SchemaName == databaseName { filteredGot = append(filteredGot, item) } } } if len(filteredGot) == 0 { got = nil } else { got = filteredGot } } if diff := cmp.Diff(tc.want, got, opts...); diff != "" { t.Errorf("Unexpected result: got %#v, want: %#v", got, tc.want) } }) } } // RunMySQLListActiveQueriesTest run tests against the mysql-list-active-queries tests func RunMySQLListActiveQueriesTest(t *testing.T, ctx context.Context, pool *sql.DB) { type queryListDetails struct { ProcessId any `json:"process_id"` Query string `json:"query"` TrxStarted any `json:"trx_started"` TrxDuration any `json:"trx_duration_seconds"` TrxWaitDuration any `json:"trx_wait_duration_seconds"` QueryTime any `json:"query_time"` TrxState string `json:"trx_state"` ProcessState string `json:"process_state"` User string `json:"user"` TrxRowsLocked any `json:"trx_rows_locked"` TrxRowsModified any `json:"trx_rows_modified"` Db string `json:"db"` } singleQueryWanted := queryListDetails{ ProcessId: any(nil), Query: "SELECT sleep(10)", TrxStarted: any(nil), TrxDuration: any(nil), TrxWaitDuration: any(nil), QueryTime: any(nil), TrxState: "", ProcessState: "User sleep", User: "", TrxRowsLocked: any(nil), TrxRowsModified: any(nil), Db: "", } invokeTcs := []struct { name string requestBody io.Reader clientSleepSecs int waitSecsBeforeCheck int wantStatusCode int want any }{ { name: "invoke list_active_queries when the system is idle", requestBody: bytes.NewBufferString(`{}`), clientSleepSecs: 0, waitSecsBeforeCheck: 0, wantStatusCode: http.StatusOK, want: []queryListDetails(nil), }, { name: "invoke list_active_queries when there is 1 ongoing but lower than the threshold", requestBody: bytes.NewBufferString(`{"min_duration_secs": 100}`), clientSleepSecs: 10, waitSecsBeforeCheck: 1, wantStatusCode: http.StatusOK, want: []queryListDetails(nil), }, { name: "invoke list_active_queries when 1 ongoing query should show up", requestBody: bytes.NewBufferString(`{"min_duration_secs": 5}`), clientSleepSecs: 0, waitSecsBeforeCheck: 5, wantStatusCode: http.StatusOK, want: []queryListDetails{singleQueryWanted}, }, { name: "invoke list_active_queries when 2 ongoing query should show up", requestBody: bytes.NewBufferString(`{"min_duration_secs": 2}`), clientSleepSecs: 10, waitSecsBeforeCheck: 3, wantStatusCode: http.StatusOK, want: []queryListDetails{singleQueryWanted, singleQueryWanted}, }, } var wg sync.WaitGroup for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { if tc.clientSleepSecs > 0 { wg.Add(1) go func() { defer wg.Done() err := pool.PingContext(ctx) if err != nil { t.Errorf("unable to connect to test database: %s", err) return } _, err = pool.ExecContext(ctx, fmt.Sprintf("SELECT sleep(%d);", tc.clientSleepSecs)) if err != nil { t.Errorf("Executing 'SELECT sleep' failed: %s", err) } }() } if tc.waitSecsBeforeCheck > 0 { time.Sleep(time.Duration(tc.waitSecsBeforeCheck) * time.Second) } const api = "http://127.0.0.1:5000/api/tool/list_active_queries/invoke" req, err := http.NewRequest(http.MethodPost, api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %v", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %v", err) } defer resp.Body.Close() if resp.StatusCode != tc.wantStatusCode { body, _ := io.ReadAll(resp.Body) t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(body)) } if tc.wantStatusCode != http.StatusOK { return } var bodyWrapper struct { Result json.RawMessage `json:"result"` } if err := json.NewDecoder(resp.Body).Decode(&bodyWrapper); err != nil { t.Fatalf("error decoding response wrapper: %v", err) } var resultString string if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil { resultString = string(bodyWrapper.Result) } var got any var details []queryListDetails if err := json.Unmarshal([]byte(resultString), &details); err != nil { t.Fatalf("failed to unmarshal nested ObjectDetails string: %v", err) } got = details if diff := cmp.Diff(tc.want, got, cmp.Comparer(func(a, b queryListDetails) bool { return a.Query == b.Query && a.ProcessState == b.ProcessState })); diff != "" { t.Errorf("Unexpected result: got %#v, want: %#v", got, tc.want) } }) } wg.Wait() } func RunMySQLListTablesMissingUniqueIndexes(t *testing.T, ctx context.Context, pool *sql.DB, databaseName string) { type listDetails struct { TableSchema string `json:"table_schema"` TableName string `json:"table_name"` } // bunch of wanted nonUniqueKeyTableName := "t03_non_unqiue_key_table" noKeyTableName := "t04_no_key_table" nonUniqueKeyTableWant := listDetails{ TableSchema: databaseName, TableName: nonUniqueKeyTableName, } noKeyTableWant := listDetails{ TableSchema: databaseName, TableName: noKeyTableName, } invokeTcs := []struct { name string requestBody io.Reader newTableName string newTablePrimaryKey bool newTableUniqueKey bool newTableNonUniqueKey bool wantStatusCode int want any }{ { name: "invoke list_tables_missing_unique_indexes when nothing to be found", requestBody: bytes.NewBufferString(`{}`), newTableName: "", newTablePrimaryKey: false, newTableUniqueKey: false, newTableNonUniqueKey: false, wantStatusCode: http.StatusOK, want: []listDetails(nil), }, { name: "invoke list_tables_missing_unique_indexes pk table will not show", requestBody: bytes.NewBufferString(`{}`), newTableName: "t01", newTablePrimaryKey: true, newTableUniqueKey: false, newTableNonUniqueKey: false, wantStatusCode: http.StatusOK, want: []listDetails(nil), }, { name: "invoke list_tables_missing_unique_indexes uk table will not show", requestBody: bytes.NewBufferString(`{}`), newTableName: "t02", newTablePrimaryKey: false, newTableUniqueKey: true, newTableNonUniqueKey: false, wantStatusCode: http.StatusOK, want: []listDetails(nil), }, { name: "invoke list_tables_missing_unique_indexes non-unique key only table will show", requestBody: bytes.NewBufferString(`{}`), newTableName: nonUniqueKeyTableName, newTablePrimaryKey: false, newTableUniqueKey: false, newTableNonUniqueKey: true, wantStatusCode: http.StatusOK, want: []listDetails{nonUniqueKeyTableWant}, }, { name: "invoke list_tables_missing_unique_indexes table with no key at all will show", requestBody: bytes.NewBufferString(`{}`), newTableName: noKeyTableName, newTablePrimaryKey: false, newTableUniqueKey: false, newTableNonUniqueKey: false, wantStatusCode: http.StatusOK, want: []listDetails{nonUniqueKeyTableWant, noKeyTableWant}, }, { name: "invoke list_tables_missing_unique_indexes table w/ both pk & uk will not show", requestBody: bytes.NewBufferString(`{}`), newTableName: "t05", newTablePrimaryKey: true, newTableUniqueKey: true, newTableNonUniqueKey: false, wantStatusCode: http.StatusOK, want: []listDetails{nonUniqueKeyTableWant, noKeyTableWant}, }, { name: "invoke list_tables_missing_unique_indexes table w/ uk & nk will not show", requestBody: bytes.NewBufferString(`{}`), newTableName: "t06", newTablePrimaryKey: false, newTableUniqueKey: true, newTableNonUniqueKey: true, wantStatusCode: http.StatusOK, want: []listDetails{nonUniqueKeyTableWant, noKeyTableWant}, }, { name: "invoke list_tables_missing_unique_indexes table w/ pk & nk will not show", requestBody: bytes.NewBufferString(`{}`), newTableName: "t07", newTablePrimaryKey: true, newTableUniqueKey: false, newTableNonUniqueKey: true, wantStatusCode: http.StatusOK, want: []listDetails{nonUniqueKeyTableWant, noKeyTableWant}, }, { name: "invoke list_tables_missing_unique_indexes with a non-exist database, nothing to show", requestBody: bytes.NewBufferString(`{"table_schema": "non-exist-database"}`), newTableName: "", newTablePrimaryKey: false, newTableUniqueKey: false, newTableNonUniqueKey: false, wantStatusCode: http.StatusOK, want: []listDetails(nil), }, { name: "invoke list_tables_missing_unique_indexes with the right database, show everything", requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_schema": "%s"}`, databaseName)), newTableName: "", newTablePrimaryKey: false, newTableUniqueKey: false, newTableNonUniqueKey: false, wantStatusCode: http.StatusOK, want: []listDetails{nonUniqueKeyTableWant, noKeyTableWant}, }, { name: "invoke list_tables_missing_unique_indexes with limited output", requestBody: bytes.NewBufferString(`{"limit": 1}`), newTableName: "", newTablePrimaryKey: false, newTableUniqueKey: false, newTableNonUniqueKey: false, wantStatusCode: http.StatusOK, want: []listDetails{nonUniqueKeyTableWant}, }, } createTableHelper := func(t *testing.T, tableName, databaseName string, primaryKey, uniqueKey, nonUniqueKey bool, ctx context.Context, pool *sql.DB) func() { var stmt strings.Builder stmt.WriteString(fmt.Sprintf("CREATE TABLE %s (", tableName)) stmt.WriteString("c1 INT") if primaryKey { stmt.WriteString(" PRIMARY KEY") } stmt.WriteString(", c2 INT, c3 CHAR(8)") if uniqueKey { stmt.WriteString(", UNIQUE(c2)") } if nonUniqueKey { stmt.WriteString(", INDEX(c3)") } stmt.WriteString(")") t.Logf("Creating table: %s", stmt.String()) if _, err := pool.ExecContext(ctx, stmt.String()); err != nil { t.Fatalf("failed executing %s: %v", stmt.String(), err) } return func() { t.Logf("Dropping table: %s", tableName) if _, err := pool.ExecContext(ctx, fmt.Sprintf("DROP TABLE %s", tableName)); err != nil { t.Errorf("failed to drop table %s: %v", tableName, err) } } } var cleanups []func() defer func() { for i := len(cleanups) - 1; i >= 0; i-- { cleanups[i]() } }() for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { if tc.newTableName != "" { cleanup := createTableHelper(t, tc.newTableName, databaseName, tc.newTablePrimaryKey, tc.newTableUniqueKey, tc.newTableNonUniqueKey, ctx, pool) cleanups = append(cleanups, cleanup) } const api = "http://127.0.0.1:5000/api/tool/list_tables_missing_unique_indexes/invoke" req, err := http.NewRequest(http.MethodPost, api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %v", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %v", err) } defer resp.Body.Close() if resp.StatusCode != tc.wantStatusCode { body, _ := io.ReadAll(resp.Body) t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(body)) } if tc.wantStatusCode != http.StatusOK { return } var bodyWrapper struct { Result json.RawMessage `json:"result"` } if err := json.NewDecoder(resp.Body).Decode(&bodyWrapper); err != nil { t.Fatalf("error decoding response wrapper: %v", err) } var resultString string if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil { resultString = string(bodyWrapper.Result) } var got any var details []listDetails if err := json.Unmarshal([]byte(resultString), &details); err != nil { t.Fatalf("failed to unmarshal nested listDetails string: %v", err) } got = details if diff := cmp.Diff(tc.want, got, cmp.Comparer(func(a, b listDetails) bool { return a.TableSchema == b.TableSchema && a.TableName == b.TableName })); diff != "" { t.Errorf("Unexpected result: got %#v, want: %#v", got, tc.want) } }) } } func RunMySQLListTableFragmentationTest(t *testing.T, databaseName, tableNameParam, tableNameAuth string) { type tableFragmentationDetails struct { TableSchema string `json:"table_schema"` TableName string `json:"table_name"` DataSize any `json:"data_size"` IndexSize any `json:"index_size"` DataFree any `json:"data_free"` FragmentationPercentage any `json:"fragmentation_percentage"` } paramTableEntryWanted := tableFragmentationDetails{ TableSchema: databaseName, TableName: tableNameParam, DataSize: any(nil), IndexSize: any(nil), DataFree: any(nil), FragmentationPercentage: any(nil), } authTableEntryWanted := tableFragmentationDetails{ TableSchema: databaseName, TableName: tableNameAuth, DataSize: any(nil), IndexSize: any(nil), DataFree: any(nil), FragmentationPercentage: any(nil), } invokeTcs := []struct { name string requestBody io.Reader wantStatusCode int want any }{ { name: "invoke list_table_fragmentation on all, no data_free threshold, expected to have 2 results", requestBody: bytes.NewBufferString(`{"data_free_threshold_bytes": 0}`), wantStatusCode: http.StatusOK, want: []tableFragmentationDetails{authTableEntryWanted, paramTableEntryWanted}, }, { name: "invoke list_table_fragmentation on all, no data_free threshold, limit to 1, expected to have 1 results", requestBody: bytes.NewBufferString(`{"data_free_threshold_bytes": 0, "limit": 1}`), wantStatusCode: http.StatusOK, want: []tableFragmentationDetails{authTableEntryWanted}, }, { name: "invoke list_table_fragmentation on all databases and 1 specific table name, no data_free threshold, expected to have 1 result", requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_name": "%s","data_free_threshold_bytes": 0}`, tableNameAuth)), wantStatusCode: http.StatusOK, want: []tableFragmentationDetails{authTableEntryWanted}, }, { name: "invoke list_table_fragmentation on 1 database and 1 specific table name, no data_free threshold, expected to have 1 result", requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_schema": "%s", "table_name": "%s", "data_free_threshold_bytes": 0}`, databaseName, tableNameParam)), wantStatusCode: http.StatusOK, want: []tableFragmentationDetails{paramTableEntryWanted}, }, { name: "invoke list_table_fragmentation on 1 database and 1 specific table name, high data_free threshold, expected to have 0 result", requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_schema": "%s", "table_name": "%s", "data_free_threshold_bytes": 1000000000}`, databaseName, tableNameParam)), wantStatusCode: http.StatusOK, want: []tableFragmentationDetails(nil), }, { name: "invoke list_table_fragmentation on 1 non-exist database, no data_free threshold, expected to have 0 result", requestBody: bytes.NewBufferString(`{"table_schema": "non_existent_database", "data_free_threshold_bytes": 0}`), wantStatusCode: http.StatusOK, want: []tableFragmentationDetails(nil), }, { name: "invoke list_table_fragmentation on 1 non-exist table, no data_free threshold, expected to have 0 result", requestBody: bytes.NewBufferString(`{"table_name": "non_existent_table", "data_free_threshold_bytes": 0}`), wantStatusCode: http.StatusOK, want: []tableFragmentationDetails(nil), }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { const api = "http://127.0.0.1:5000/api/tool/list_table_fragmentation/invoke" req, err := http.NewRequest(http.MethodPost, api, tc.requestBody) if err != nil { t.Fatalf("unable to create request: %v", err) } req.Header.Add("Content-type", "application/json") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %v", err) } defer resp.Body.Close() if resp.StatusCode != tc.wantStatusCode { body, _ := io.ReadAll(resp.Body) t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(body)) } if tc.wantStatusCode != http.StatusOK { return } var bodyWrapper struct { Result json.RawMessage `json:"result"` } if err := json.NewDecoder(resp.Body).Decode(&bodyWrapper); err != nil { t.Fatalf("error decoding response wrapper: %v", err) } var resultString string if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil { resultString = string(bodyWrapper.Result) } var got any var details []tableFragmentationDetails if err := json.Unmarshal([]byte(resultString), &details); err != nil { t.Fatalf("failed to unmarshal outer JSON array into []tableInfo: %v", err) } got = details if diff := cmp.Diff(tc.want, got, cmp.Comparer(func(a, b tableFragmentationDetails) bool { return a.TableSchema == b.TableSchema && a.TableName == b.TableName })); diff != "" { t.Errorf("Unexpected result: got %#v, want: %#v", got, tc.want) } }) } } // RunMSSQLListTablesTest run tests againsts the mssql-list-tables tools. func RunMSSQLListTablesTest(t *testing.T, tableNameParam, tableNameAuth string) { // TableNameParam columns to construct want. const paramTableColumns = `[ {"column_name": "id", "data_type": "INT", "column_ordinal_position": 1, "is_not_nullable": true}, {"column_name": "name", "data_type": "VARCHAR(255)", "column_ordinal_position": 2, "is_not_nullable": false} ]` // TableNameAuth columns to construct want const authTableColumns = `[ {"column_name": "id", "data_type": "INT", "column_ordinal_position": 1, "is_not_nullable": true}, {"column_name": "name", "data_type": "VARCHAR(255)", "column_ordinal_position": 2, "is_not_nullable": false}, {"column_name": "email", "data_type": "VARCHAR(255)", "column_ordinal_position": 3, "is_not_nullable": false} ]` const ( // Template to construct detailed output want. detailedObjectTemplate = `{ "schema_name": "dbo", "object_name": "%[1]s", "object_details": { "owner": "dbo", "triggers": [], "columns": %[2]s, "object_name": "%[1]s", "object_type": "TABLE", "schema_name": "dbo" } }` // Template to construct simple output want simpleObjectTemplate = `{"object_name":"%s", "schema_name":"dbo", "object_details":{"name":"%s"}}` ) // Helper to build json for detailed want getDetailedWant := func(tableName, columnJSON string) string { return fmt.Sprintf(detailedObjectTemplate, tableName, columnJSON) } // Helper to build template for simple want getSimpleWant := func(tableName string) string { return fmt.Sprintf(simpleObjectTemplate, tableName, tableName) } invokeTcs := []struct { name string api string requestBody string wantStatusCode int want string isAllTables bool }{ { name: "invoke list_tables for all tables detailed output", api: "http://127.0.0.1:5000/api/tool/list_tables/invoke", requestBody: `{"table_names": ""}`, wantStatusCode: http.StatusOK, want: fmt.Sprintf("[%s,%s]", getDetailedWant(tableNameAuth, authTableColumns), getDetailedWant(tableNameParam, paramTableColumns)), isAllTables: true, }, { name: "invoke list_tables for all tables simple output", api: "http://127.0.0.1:5000/api/tool/list_tables/invoke", requestBody: `{"table_names": "", "output_format": "simple"}`, wantStatusCode: http.StatusOK, want: fmt.Sprintf("[%s,%s]", getSimpleWant(tableNameAuth), getSimpleWant(tableNameParam)), isAllTables: true, }, { name: "invoke list_tables detailed output", api: "http://127.0.0.1:5000/api/tool/list_tables/invoke", requestBody: fmt.Sprintf(`{"table_names": "%s"}`, tableNameAuth), wantStatusCode: http.StatusOK, want: fmt.Sprintf("[%s]", getDetailedWant(tableNameAuth, authTableColumns)), }, { name: "invoke list_tables simple output", api: "http://127.0.0.1:5000/api/tool/list_tables/invoke", requestBody: fmt.Sprintf(`{"table_names": "%s", "output_format": "simple"}`, tableNameAuth), wantStatusCode: http.StatusOK, want: fmt.Sprintf("[%s]", getSimpleWant(tableNameAuth)), }, { name: "invoke list_tables with invalid output format", api: "http://127.0.0.1:5000/api/tool/list_tables/invoke", requestBody: `{"table_names": "", "output_format": "abcd"}`, wantStatusCode: http.StatusBadRequest, }, { name: "invoke list_tables with malformed table_names parameter", api: "http://127.0.0.1:5000/api/tool/list_tables/invoke", requestBody: `{"table_names": 12345, "output_format": "detailed"}`, wantStatusCode: http.StatusBadRequest, }, { name: "invoke list_tables with multiple table names", api: "http://127.0.0.1:5000/api/tool/list_tables/invoke", requestBody: fmt.Sprintf(`{"table_names": "%s,%s"}`, tableNameParam, tableNameAuth), wantStatusCode: http.StatusOK, want: fmt.Sprintf("[%s,%s]", getDetailedWant(tableNameAuth, authTableColumns), getDetailedWant(tableNameParam, paramTableColumns)), }, { name: "invoke list_tables with non-existent table", api: "http://127.0.0.1:5000/api/tool/list_tables/invoke", requestBody: `{"table_names": "non_existent_table"}`, wantStatusCode: http.StatusOK, want: `null`, }, { name: "invoke list_tables with one existing and one non-existent table", api: "http://127.0.0.1:5000/api/tool/list_tables/invoke", requestBody: fmt.Sprintf(`{"table_names": "%s,non_existent_table"}`, tableNameParam), wantStatusCode: http.StatusOK, want: fmt.Sprintf("[%s]", getDetailedWant(tableNameParam, paramTableColumns)), }, } for _, tc := range invokeTcs { t.Run(tc.name, func(t *testing.T) { resp, respBytes := RunRequest(t, http.MethodPost, tc.api, bytes.NewBuffer([]byte(tc.requestBody)), nil) if resp.StatusCode != tc.wantStatusCode { t.Fatalf("response status code is not %d, got %d: %s", tc.wantStatusCode, resp.StatusCode, string(respBytes)) } if tc.wantStatusCode == http.StatusOK { var bodyWrapper map[string]json.RawMessage if err := json.Unmarshal(respBytes, &bodyWrapper); err != nil { t.Fatalf("error parsing response wrapper: %s, body: %s", err, string(respBytes)) } resultJSON, ok := bodyWrapper["result"] if !ok { t.Fatal("unable to find 'result' in response body") } var resultString string if err := json.Unmarshal(resultJSON, &resultString); err != nil { if string(resultJSON) == "null" { resultString = "null" } else { t.Fatalf("'result' is not a JSON-encoded string: %s", err) } } var got, want []any if err := json.Unmarshal([]byte(resultString), &got); err != nil { t.Fatalf("failed to unmarshal actual result string: %v", err) } if err := json.Unmarshal([]byte(tc.want), &want); err != nil { t.Fatalf("failed to unmarshal expected want string: %v", err) } for _, item := range got { itemMap, ok := item.(map[string]any) if !ok { continue } detailsStr, ok := itemMap["object_details"].(string) if !ok { continue } var detailsMap map[string]any if err := json.Unmarshal([]byte(detailsStr), &detailsMap); err != nil { t.Fatalf("failed to unmarshal nested object_details string: %v", err) } // clean unpredictable fields delete(detailsMap, "constraints") delete(detailsMap, "indexes") itemMap["object_details"] = detailsMap } // Checking only the default dbo schema where the test tables are created to avoid brittle tests. if tc.isAllTables { var filteredGot []any for _, item := range got { if tableMap, ok := item.(map[string]interface{}); ok { if schema, ok := tableMap["schema_name"]; ok && schema == "dbo" { filteredGot = append(filteredGot, item) } } } got = filteredGot } sort.SliceStable(got, func(i, j int) bool { return fmt.Sprintf("%v", got[i]) < fmt.Sprintf("%v", got[j]) }) sort.SliceStable(want, func(i, j int) bool { return fmt.Sprintf("%v", want[i]) < fmt.Sprintf("%v", want[j]) }) if !reflect.DeepEqual(got, want) { gotJSON, _ := json.MarshalIndent(got, "", " ") wantJSON, _ := json.MarshalIndent(want, "", " ") t.Errorf("Unexpected result:\ngot:\n%s\n\nwant:\n%s", string(gotJSON), string(wantJSON)) } } }) } } // RunRequest is a helper function to send HTTP requests and return the response func RunRequest(t *testing.T, method, url string, body io.Reader, headers map[string]string) (*http.Response, []byte) { // Send request req, err := http.NewRequest(method, url, body) if err != nil { t.Fatalf("unable to create request: %s", err) } req.Header.Set("Content-type", "application/json") for k, v := range headers { req.Header.Set(k, v) } resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("unable to send request: %s", err) } respBody, err := io.ReadAll(resp.Body) if err != nil { t.Fatalf("unable to read request body: %s", err) } defer resp.Body.Close() return resp, respBody } ```