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

Code before:
			}
			bodyBytes, err := json.Marshal(requestBodyMap)
			if err != nil {
				t.Fatalf("failed to marshal request body: %v", err)
			}
			body := bytes.NewBuffer(bodyBytes)

			req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1:5000/api/tool/analyze-contribution-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))
			}

			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 err := json.Unmarshal(bodyBytes, &respBody); err != nil {
```

Code after:

			if tc.wantInError != "" {
				bodyBytes, _ := io.ReadAll(resp.Body)
				if !strings.Contains(string(bodyBytes), tc.wantInError) {
					t.Errorf("unexpected error message: got %q, want to contain %q", string(bodyBytes), tc.wantInError)
				}
			}
		})
	}
}

func runAnalyzeContributionWithRestriction(t *testing.T, allowedTableFullName, disallowedTableFullName string) {
	allowedTableUnquoted := strings.ReplaceAll(allowedTableFullName, "`", "")
	disallowedTableUnquoted := strings.ReplaceAll(disallowedTableFullName, "`", "")
	disallowedDatasetFQN := strings.Join(strings.Split(disallowedTableUnquoted, ".")[0:2], ".")

	testCases := []struct {
		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) {