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

Code before:

			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))
			}

			respBodyBytes, _ := io.ReadAll(resp.Body)
			var respBody map[string]interface{}
			if err := json.Unmarshal(respBodyBytes, &respBody); err != nil {
				t.Fatalf("error parsing response body: %v", err)
			}

			if tc.wantInResult != "" {
				got, ok := respBody["result"].(string)
				if !ok {
					t.Fatalf("unable to find result in response body")
				}

				if !strings.Contains(got, tc.wantInResult) {
					t.Errorf("unexpected result: got %q, want to contain %q", string(respBodyBytes), tc.wantInResult)
				}
			}

			if tc.wantInError != "" {
				if !strings.Contains(string(respBodyBytes), tc.wantInError) {
					t.Errorf("unexpected error message: got %q, want to contain %q", string(respBodyBytes), tc.wantInError)
				}
			}
		})
	}
}


Review comment (suggestion):
```suggestion
				if !strings.Contains(string(bodyBytes), tc.wantInError) {
					t.Errorf("unexpected error message: got %q, want to contain %q", string(bodyBytes), tc.wantInError)
```

Code after:
		name           string
		inputData      string
		wantStatusCode int
		wantInResult   string
		wantInError    string
	}{
		{
			name:           "invoke with allowed table name",
			inputData:      allowedTableUnquoted,
			wantStatusCode: http.StatusOK,
			wantInResult:   `"relative_difference"`,
		},
		{
			name:           "invoke with disallowed table name",
			inputData:      disallowedTableUnquoted,
			wantStatusCode: http.StatusBadRequest,
			wantInError:    fmt.Sprintf("access to dataset '%s' (from table '%s') is not allowed", disallowedDatasetFQN, disallowedTableUnquoted),
		},
		{
			name:           "invoke with query on allowed table",
			inputData:      fmt.Sprintf("SELECT * FROM %s", allowedTableFullName),
			wantStatusCode: http.StatusOK,
			wantInResult:   `"relative_difference"`,
		},
		{
			name:           "invoke with query on disallowed table",
			inputData:      fmt.Sprintf("SELECT * FROM %s", disallowedTableFullName),
			wantStatusCode: http.StatusBadRequest,
			wantInError:    fmt.Sprintf("query in input_data accesses dataset '%s', which is not in the allowed list", disallowedDatasetFQN),
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			requestBodyMap := map[string]any{
				"input_data":          tc.inputData,
				"contribution_metric": "SUM(metric)",
				"is_test_col":         "is_test",
				"dimension_id_cols":   []string{"dim1", "dim2"},
			}
			bodyBytes, err := json.Marshal(requestBodyMap)
			if err != nil {
				t.Fatalf("failed to marshal request body: %v", err)
			}
			body := bytes.NewBuffer(bodyBytes)

			resp, bodyBytes := tests.RunRequest(t, http.MethodPost, "http://127.0.0.1:5000/api/tool/analyze-contribution-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))
			}