Repository: googleapis/genai-toolbox
File: tests/bigquery/bigquery_integration_test.go (Go)

Code before:
func runListDatasetIdsWithRestriction(t *testing.T, allowedDatasetName1, allowedDatasetName2 string) {
	testCases := []struct {
		name           string
		wantStatusCode int
		wantResult     string
	}{
		{
			name:           "invoke list-dataset-ids with restriction",
			wantStatusCode: http.StatusOK,
			wantResult:     fmt.Sprintf(`["%s.%s","%s.%s"]`, BigqueryProject, allowedDatasetName1, BigqueryProject, allowedDatasetName2),
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			body := bytes.NewBuffer([]byte(`{}`))
			req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1:5000/api/tool/list-dataset-ids-restricted/invoke", body)
			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 != tc.wantStatusCode {
				bodyBytes, _ := io.ReadAll(resp.Body)
				t.Fatalf("unexpected status code: got %d, want %d. Body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
			}

			var respBody map[string]interface{}
			if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
				t.Fatalf("error parsing response body: %v", err)
			}
			got, ok := respBody["result"].(string)
			if !ok {
				t.Fatalf("unable to find result in response body")
			}
			if got != tc.wantResult {
				t.Errorf("unexpected result: got %q, want %q", got, tc.wantResult)
			}
		})
	}
}

func runListTableIdsWithRestriction(t *testing.T, allowedDatasetName, disallowedDatasetName string, allowedTableNames ...string) {
	sort.Strings(allowedTableNames)
	var quotedNames []string
	for _, name := range allowedTableNames {

Review comment (suggestion):
```suggestion
			resp, bodyBytes, err := tests.RunRequest(t, http.MethodPost, "http://127.0.0.1:5000/api/tool/list-dataset-ids-restricted/invoke", body, nil)
			if err != nil {
				t.Fatalf("unable to create request: %s", err)
			}
```

Code after:
func runListDatasetIdsWithRestriction(t *testing.T, allowedDatasetName1, allowedDatasetName2 string) {
	testCases := []struct {
		name           string
		wantStatusCode int
		wantResult     string
	}{
		{
			name:           "invoke list-dataset-ids with restriction",
			wantStatusCode: http.StatusOK,
			wantResult:     fmt.Sprintf(`["%s.%s","%s.%s"]`, BigqueryProject, allowedDatasetName1, BigqueryProject, allowedDatasetName2),
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			body := bytes.NewBuffer([]byte(`{}`))
			resp, bodyBytes := tests.RunRequest(t, http.MethodPost, "http://127.0.0.1:5000/api/tool/list-dataset-ids-restricted/invoke", body, nil)

			if resp.StatusCode != tc.wantStatusCode {
				t.Fatalf("unexpected status code: got %d, want %d. Body: %s", resp.StatusCode, tc.wantStatusCode, string(bodyBytes))
			}

			var respBody map[string]interface{}
			if err := json.Unmarshal(bodyBytes, &respBody); err != nil {
				t.Fatalf("error parsing response body: %v", err)
			}
			got, ok := respBody["result"].(string)
			if !ok {
				t.Fatalf("unable to find result in response body")
			}
			if got != tc.wantResult {
				t.Errorf("unexpected result: got %q, want %q", got, tc.wantResult)
			}
		})
	}
}

func runListTableIdsWithRestriction(t *testing.T, allowedDatasetName, disallowedDatasetName string, allowedTableNames ...string) {
	sort.Strings(allowedTableNames)
	var quotedNames []string
	for _, name := range allowedTableNames {
		quotedNames = append(quotedNames, fmt.Sprintf(`"%s"`, name))
	}
	wantResult := fmt.Sprintf(`[%s]`, strings.Join(quotedNames, ","))

	testCases := []struct {
		name           string
		dataset        string
		wantStatusCode int
		wantInResult   string
		wantInError    string