wehub-resource-sync
e04ed9c211
CF: Deploy Dev Docs / deploy (push) Has been cancelled
Sync Labels / build (push) Has been cancelled
tests / unit tests (macos-latest) (push) Has been cancelled
tests / unit tests (windows-latest) (push) Has been cancelled
tests / unit tests (ubuntu-latest) (push) Has been cancelled
907 lines
26 KiB
Go
907 lines
26 KiB
Go
// Copyright 2026 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 dataproc
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"reflect"
|
|
"regexp"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
dataproc "cloud.google.com/go/dataproc/v2/apiv1"
|
|
"cloud.google.com/go/dataproc/v2/apiv1/dataprocpb"
|
|
"github.com/google/go-cmp/cmp"
|
|
dataprocsrc "github.com/googleapis/mcp-toolbox/internal/sources/dataproc"
|
|
"github.com/googleapis/mcp-toolbox/internal/testutils"
|
|
"github.com/googleapis/mcp-toolbox/tests"
|
|
"google.golang.org/api/iterator"
|
|
"google.golang.org/api/option"
|
|
"google.golang.org/protobuf/encoding/protojson"
|
|
"google.golang.org/protobuf/testing/protocmp"
|
|
)
|
|
|
|
var (
|
|
dataprocRegion = os.Getenv("DATAPROC_REGION")
|
|
dataprocProject = os.Getenv("DATAPROC_PROJECT")
|
|
|
|
// dataprocListJobsCluster is the name of a cluster in the project that has jobs.
|
|
//
|
|
// This is necessary to work around a performance issue in the Dataproc API where listing all
|
|
// jobs in a project is very slow.
|
|
dataprocListJobsCluster = os.Getenv("DATAPROC_LIST_JOBS_CLUSTER")
|
|
)
|
|
|
|
const (
|
|
clusterURLPrefix = "https://console.cloud.google.com/dataproc/clusters/"
|
|
jobURLPrefix = "https://console.cloud.google.com/dataproc/jobs/"
|
|
logsURLPrefix = "https://console.cloud.google.com/logs/viewer?"
|
|
)
|
|
|
|
func getDataprocVars(t *testing.T) map[string]any {
|
|
switch "" {
|
|
case dataprocRegion:
|
|
t.Fatal("'DATAPROC_REGION' not set")
|
|
case dataprocProject:
|
|
t.Fatal("'DATAPROC_PROJECT' not set")
|
|
case dataprocListJobsCluster:
|
|
t.Fatal("'DATAPROC_LIST_JOBS_CLUSTER' not set")
|
|
}
|
|
|
|
return map[string]any{
|
|
"type": "dataproc",
|
|
"project": dataprocProject,
|
|
"region": dataprocRegion,
|
|
}
|
|
}
|
|
|
|
func TestDataprocClustersToolEndpoints(t *testing.T) {
|
|
sourceConfig := getDataprocVars(t)
|
|
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) // Clusters take time
|
|
defer cancel()
|
|
|
|
toolsFile := map[string]any{
|
|
"sources": map[string]any{
|
|
"my-dataproc": sourceConfig,
|
|
},
|
|
"authServices": map[string]any{
|
|
"my-google-auth": map[string]any{
|
|
"type": "google",
|
|
"clientId": tests.ClientId,
|
|
},
|
|
},
|
|
"tools": map[string]any{
|
|
"get-cluster": map[string]any{
|
|
"type": "dataproc-get-cluster",
|
|
"source": "my-dataproc",
|
|
},
|
|
"get-cluster-with-auth": map[string]any{
|
|
"type": "dataproc-get-cluster",
|
|
"source": "my-dataproc",
|
|
"authRequired": []string{"my-google-auth"},
|
|
},
|
|
"get-job": map[string]any{
|
|
"type": "dataproc-get-job",
|
|
"source": "my-dataproc",
|
|
},
|
|
"get-job-with-auth": map[string]any{
|
|
"type": "dataproc-get-job",
|
|
"source": "my-dataproc",
|
|
"authRequired": []string{"my-google-auth"},
|
|
},
|
|
"list-clusters": map[string]any{
|
|
"type": "dataproc-list-clusters",
|
|
"source": "my-dataproc",
|
|
},
|
|
"list-clusters-with-auth": map[string]any{
|
|
"type": "dataproc-list-clusters",
|
|
"source": "my-dataproc",
|
|
"authRequired": []string{"my-google-auth"},
|
|
},
|
|
"list-jobs": map[string]any{
|
|
"type": "dataproc-list-jobs",
|
|
"source": "my-dataproc",
|
|
},
|
|
"list-jobs-with-auth": map[string]any{
|
|
"type": "dataproc-list-jobs",
|
|
"source": "my-dataproc",
|
|
"authRequired": []string{"my-google-auth"},
|
|
},
|
|
},
|
|
}
|
|
|
|
args := []string{"--enable-api"}
|
|
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)
|
|
}
|
|
|
|
endpoint := fmt.Sprintf("%s-dataproc.googleapis.com:443", dataprocRegion)
|
|
clusterClient, err := dataproc.NewClusterControllerClient(ctx, option.WithEndpoint(endpoint))
|
|
if err != nil {
|
|
t.Fatalf("failed to create dataproc client: %v", err)
|
|
}
|
|
defer clusterClient.Close()
|
|
|
|
jobClient, err := dataproc.NewJobControllerClient(ctx, option.WithEndpoint(endpoint))
|
|
if err != nil {
|
|
t.Fatalf("failed to create dataproc client: %v", err)
|
|
}
|
|
defer jobClient.Close()
|
|
|
|
t.Run("get-cluster", func(t *testing.T) {
|
|
clusterName := listClustersRpc(t, clusterClient, ctx, "", 1)[0].Name
|
|
t.Run("success", func(t *testing.T) {
|
|
t.Parallel()
|
|
runGetClusterTest(t, clusterClient, ctx, clusterName)
|
|
})
|
|
t.Run("errors", func(t *testing.T) {
|
|
t.Parallel()
|
|
missingClusterFullName := fmt.Sprintf("projects/%s/regions/%s/clusters/INVALID_CLUSTER", dataprocProject, dataprocRegion)
|
|
tcs := []struct {
|
|
name string
|
|
toolName string
|
|
request map[string]any
|
|
wantCode int
|
|
wantMsg string
|
|
}{
|
|
{
|
|
name: "missing cluster",
|
|
toolName: "get-cluster",
|
|
request: map[string]any{"clusterName": "INVALID_CLUSTER"},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: fmt.Sprintf("Not found: Cluster projects/%s/regions/%s/clusters/INVALID_CLUSTER", dataprocProject, dataprocRegion),
|
|
},
|
|
{
|
|
name: "full cluster name",
|
|
toolName: "get-cluster",
|
|
request: map[string]any{"clusterName": missingClusterFullName},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: fmt.Sprintf("clusterName must be a short name without '/': %s", missingClusterFullName),
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
testError(t, tc.toolName, tc.request, tc.wantCode, tc.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
t.Run("auth", func(t *testing.T) {
|
|
t.Parallel()
|
|
runAuthTest(t, "get-cluster-with-auth", map[string]any{"clusterName": shortName(clusterName)}, http.StatusOK)
|
|
})
|
|
})
|
|
|
|
t.Run("get-job", func(t *testing.T) {
|
|
jobId := listJobsRpc(t, jobClient, ctx, "", 1)[0].ID
|
|
t.Run("success", func(t *testing.T) {
|
|
t.Parallel()
|
|
runGetJobTest(t, jobClient, ctx, jobId)
|
|
})
|
|
t.Run("errors", func(t *testing.T) {
|
|
t.Parallel()
|
|
missingJobFullName := fmt.Sprintf("projects/%s/regions/%s/jobs/INVALID_JOB", dataprocProject, dataprocRegion)
|
|
tcs := []struct {
|
|
name string
|
|
toolName string
|
|
request map[string]any
|
|
wantCode int
|
|
wantMsg string
|
|
}{
|
|
{
|
|
name: "missing job",
|
|
toolName: "get-job",
|
|
request: map[string]any{"jobId": "INVALID_JOB"},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: fmt.Sprintf("Not found: Job projects/%s/regions/%s/jobs/INVALID_JOB", dataprocProject, dataprocRegion),
|
|
},
|
|
{
|
|
name: "full job name",
|
|
toolName: "get-job",
|
|
request: map[string]any{"jobId": missingJobFullName},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: fmt.Sprintf("jobId must be a short name without '/': %s", missingJobFullName),
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
testError(t, tc.toolName, tc.request, tc.wantCode, tc.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
t.Run("auth", func(t *testing.T) {
|
|
t.Parallel()
|
|
runAuthTest(t, "get-job-with-auth", map[string]any{"jobId": jobId}, http.StatusOK)
|
|
})
|
|
})
|
|
t.Run("list-clusters", func(t *testing.T) {
|
|
t.Run("success", func(t *testing.T) {
|
|
runListClustersTest(t, clusterClient, ctx)
|
|
})
|
|
t.Run("errors", func(t *testing.T) {
|
|
t.Parallel()
|
|
tcs := []struct {
|
|
name string
|
|
toolName string
|
|
request map[string]any
|
|
wantCode int
|
|
wantMsg string
|
|
}{
|
|
{
|
|
name: "zero page size",
|
|
toolName: "list-clusters",
|
|
request: map[string]any{"pageSize": 0},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: "pageSize must be positive: 0",
|
|
},
|
|
{
|
|
name: "negative page size",
|
|
toolName: "list-clusters",
|
|
request: map[string]any{"pageSize": -1},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: "pageSize must be positive: -1",
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
testError(t, tc.toolName, tc.request, tc.wantCode, tc.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
t.Run("auth", func(t *testing.T) {
|
|
t.Parallel()
|
|
runAuthTest(t, "list-clusters-with-auth", map[string]any{"pageSize": 1}, http.StatusOK)
|
|
})
|
|
})
|
|
|
|
t.Run("list-jobs", func(t *testing.T) {
|
|
t.Run("success", func(t *testing.T) {
|
|
runListJobsTest(t, jobClient, ctx)
|
|
})
|
|
t.Run("errors", func(t *testing.T) {
|
|
t.Parallel()
|
|
tcs := []struct {
|
|
name string
|
|
toolName string
|
|
request map[string]any
|
|
wantCode int
|
|
wantMsg string
|
|
}{
|
|
{
|
|
name: "zero page size",
|
|
toolName: "list-jobs",
|
|
request: map[string]any{"pageSize": 0},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: "pageSize must be positive: 0",
|
|
},
|
|
{
|
|
name: "negative page size",
|
|
toolName: "list-jobs",
|
|
request: map[string]any{"pageSize": -1},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: "pageSize must be positive: -1",
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
testError(t, tc.toolName, tc.request, tc.wantCode, tc.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
t.Run("auth", func(t *testing.T) {
|
|
t.Parallel()
|
|
runAuthTest(t, "list-jobs-with-auth", map[string]any{
|
|
"pageSize": 1,
|
|
"filter": "clusterName = " + dataprocListJobsCluster,
|
|
}, http.StatusOK)
|
|
})
|
|
})
|
|
|
|
}
|
|
|
|
func invokeTool(toolName string, request map[string]any, headers map[string]string) (*http.Response, error) {
|
|
requestBytes, err := json.Marshal(request)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
|
}
|
|
|
|
url := fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/invoke", toolName)
|
|
req, err := http.NewRequest(http.MethodPost, url, bytes.NewBuffer(requestBytes))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("unable to create request: %w", err)
|
|
}
|
|
req.Header.Add("Content-type", "application/json")
|
|
for k, v := range headers {
|
|
req.Header.Add(k, v)
|
|
}
|
|
|
|
return http.DefaultClient.Do(req)
|
|
}
|
|
|
|
func runListClustersTest(t *testing.T, client *dataproc.ClusterControllerClient, ctx context.Context) {
|
|
tcs := []struct {
|
|
name string
|
|
filter string
|
|
pageSize int
|
|
numPages int
|
|
wantN int
|
|
}{
|
|
{name: "one page", pageSize: 2, numPages: 1, wantN: 2},
|
|
{name: "two pages", pageSize: 1, numPages: 2, wantN: 2},
|
|
{name: "5 clusters", pageSize: 5, numPages: 1, wantN: 5},
|
|
{name: "omit page size", numPages: 1, wantN: 20},
|
|
{
|
|
name: "filtered",
|
|
filter: "status.state = STOPPED",
|
|
pageSize: 2,
|
|
numPages: 1,
|
|
wantN: 2,
|
|
},
|
|
{
|
|
name: "empty",
|
|
filter: "status.state = STOPPED AND status.state = RUNNING",
|
|
pageSize: 1,
|
|
numPages: 1,
|
|
wantN: 0,
|
|
},
|
|
}
|
|
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
want := []dataprocsrc.Cluster{}
|
|
if tc.wantN > 0 {
|
|
want = listClustersRpc(t, client, ctx, tc.filter, tc.wantN)
|
|
}
|
|
|
|
actual := []dataprocsrc.Cluster{}
|
|
var pageToken string
|
|
|
|
for i := 0; i < tc.numPages; i++ {
|
|
request := map[string]any{
|
|
"filter": tc.filter,
|
|
"pageToken": pageToken,
|
|
}
|
|
if tc.pageSize > 0 {
|
|
request["pageSize"] = tc.pageSize
|
|
}
|
|
|
|
resp, err := invokeTool("list-clusters", request, nil)
|
|
if err != nil {
|
|
t.Fatalf("invokeTool failed: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
bodyBytes, _ := io.ReadAll(resp.Body)
|
|
t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
|
|
}
|
|
|
|
var body map[string]any
|
|
if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
|
|
t.Fatalf("error parsing response body: %v", err)
|
|
}
|
|
|
|
result, ok := body["result"].(string)
|
|
if !ok {
|
|
t.Fatalf("unable to find result in response body")
|
|
}
|
|
|
|
var listResponse dataprocsrc.ListClustersResponse
|
|
if err := json.Unmarshal([]byte(result), &listResponse); err != nil {
|
|
t.Fatalf("error unmarshalling result: %s", err)
|
|
}
|
|
actual = append(actual, listResponse.Clusters...)
|
|
pageToken = listResponse.NextPageToken
|
|
}
|
|
|
|
if !reflect.DeepEqual(actual, want) {
|
|
t.Fatalf("unexpected clusters: got %+v, want %+v", actual, want)
|
|
}
|
|
|
|
// want has URLs because it's created from Batch instances by the same utility function
|
|
// used by the tool internals. Double-check that the URLs are reasonable.
|
|
for _, cluster := range want {
|
|
if !strings.HasPrefix(cluster.ConsoleURL, clusterURLPrefix) {
|
|
t.Errorf("unexpected consoleUrl in cluster: %#v", cluster)
|
|
}
|
|
if !strings.HasPrefix(cluster.LogsURL, logsURLPrefix) {
|
|
t.Errorf("unexpected logsUrl in cluster: %#v", cluster)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func runGetClusterTest(t *testing.T, client *dataproc.ClusterControllerClient, ctx context.Context, fullName string) {
|
|
// First get the cluster details directly from the Go proto API.
|
|
req := &dataprocpb.GetClusterRequest{
|
|
ProjectId: dataprocProject,
|
|
Region: dataprocRegion,
|
|
ClusterName: fullName[strings.LastIndex(fullName, "/")+1:],
|
|
}
|
|
rawWantClusterPb, err := client.GetCluster(ctx, req)
|
|
if err != nil {
|
|
t.Fatalf("failed to get cluster: %s", err)
|
|
}
|
|
|
|
// Trim unknown fields from the proto by marshalling and unmarshalling.
|
|
jsonBytes, err := protojson.Marshal(rawWantClusterPb)
|
|
if err != nil {
|
|
t.Fatalf("failed to marshal cluster to JSON: %s", err)
|
|
}
|
|
var wantClusterPb dataprocpb.Cluster
|
|
if err := protojson.Unmarshal(jsonBytes, &wantClusterPb); err != nil {
|
|
t.Fatalf("error unmarshalling result: %s", err)
|
|
}
|
|
|
|
shortName := fullName[strings.LastIndex(fullName, "/")+1:]
|
|
|
|
tcs := []struct {
|
|
name string
|
|
clusterName string
|
|
want *dataprocpb.Cluster
|
|
}{
|
|
{
|
|
name: "found cluster",
|
|
clusterName: shortName,
|
|
want: &wantClusterPb,
|
|
},
|
|
}
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
request := map[string]any{"clusterName": tc.clusterName}
|
|
resp, err := invokeTool("get-cluster", request, nil)
|
|
if err != nil {
|
|
t.Fatalf("invokeTool failed: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
t.Fatalf("status code got %d, want 200. Body: %s", resp.StatusCode, body)
|
|
}
|
|
|
|
var body map[string]any
|
|
if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
|
|
t.Fatalf("error parsing response body: %v", err)
|
|
}
|
|
resultStr, ok := body["result"].(string)
|
|
if !ok {
|
|
t.Fatalf("result is not a string, got %T", body["result"])
|
|
}
|
|
var wrappedResult map[string]any
|
|
if err := json.Unmarshal([]byte(resultStr), &wrappedResult); err != nil {
|
|
t.Fatalf("error unmarshalling result: %s", err)
|
|
}
|
|
|
|
consoleURL, ok := wrappedResult["consoleUrl"].(string)
|
|
if !ok || !strings.HasPrefix(consoleURL, clusterURLPrefix) {
|
|
t.Errorf("unexpected consoleUrl: %v", consoleURL)
|
|
}
|
|
logsURL, ok := wrappedResult["logsUrl"].(string)
|
|
if !ok || !strings.HasPrefix(logsURL, logsURLPrefix) {
|
|
t.Errorf("unexpected logsUrl: %v", logsURL)
|
|
}
|
|
|
|
clusterJSON, err := json.Marshal(wrappedResult["cluster"])
|
|
if err != nil {
|
|
t.Fatalf("failed to marshal cluster: %v", err)
|
|
}
|
|
|
|
// Unmarshal JSON to proto for proto-aware deep comparison.
|
|
var cluster dataprocpb.Cluster
|
|
if err := protojson.Unmarshal(clusterJSON, &cluster); err != nil {
|
|
t.Fatalf("error unmarshalling cluster from wrapped result: %s", err)
|
|
}
|
|
|
|
if !cmp.Equal(&cluster, tc.want, protocmp.Transform()) {
|
|
diff := cmp.Diff(&cluster, tc.want, protocmp.Transform())
|
|
t.Errorf("GetCluster() returned diff (-got +want):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("errors", func(t *testing.T) {
|
|
tcs := []struct {
|
|
name string
|
|
request map[string]any
|
|
wantCode int
|
|
wantMsg string
|
|
}{
|
|
{
|
|
name: "missing clusterName",
|
|
request: map[string]any{},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: "missing required parameter: clusterName",
|
|
},
|
|
{
|
|
name: "invalid name with slash",
|
|
request: map[string]any{"clusterName": "projects/foo/regions/bar/clusters/baz"}, // Full name requires matching project/region
|
|
wantCode: http.StatusOK,
|
|
wantMsg: "clusterName must be a short name without '/'",
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
testError(t, "get-cluster", tc.request, tc.wantCode, tc.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
}
|
|
|
|
func listClustersRpc(t *testing.T, client *dataproc.ClusterControllerClient, ctx context.Context, filter string, n int) []dataprocsrc.Cluster {
|
|
req := &dataprocpb.ListClustersRequest{
|
|
ProjectId: dataprocProject,
|
|
Region: dataprocRegion,
|
|
PageSize: int32(n),
|
|
}
|
|
if filter != "" {
|
|
req.Filter = filter
|
|
}
|
|
|
|
it := client.ListClusters(ctx, req)
|
|
pager := iterator.NewPager(it, n, "")
|
|
var clusterPbs []*dataprocpb.Cluster
|
|
_, err := pager.NextPage(&clusterPbs)
|
|
if err != nil {
|
|
t.Fatalf("failed to list clusters: %s", err)
|
|
}
|
|
|
|
clusters, err := dataprocsrc.ToClusters(clusterPbs, dataprocRegion)
|
|
if err != nil {
|
|
t.Fatalf("failed to convert clusters to JSON: %v", err)
|
|
}
|
|
|
|
return clusters
|
|
}
|
|
|
|
func runAuthTest(t *testing.T, toolName string, request map[string]any, wantStatus int) {
|
|
idToken, err := tests.GetGoogleIdToken(t)
|
|
if err != nil {
|
|
t.Fatalf("error getting Google ID token: %s", err)
|
|
}
|
|
|
|
tcs := []struct {
|
|
name string
|
|
headers map[string]string
|
|
wantCode int
|
|
}{
|
|
{
|
|
name: "valid token",
|
|
headers: map[string]string{"my-google-auth_token": idToken},
|
|
wantCode: wantStatus,
|
|
},
|
|
{
|
|
name: "invalid token",
|
|
headers: map[string]string{"my-google-auth_token": "INVALID"},
|
|
wantCode: http.StatusUnauthorized,
|
|
},
|
|
{
|
|
name: "missing header",
|
|
headers: nil,
|
|
wantCode: http.StatusUnauthorized,
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
resp, err := invokeTool(toolName, request, tc.headers)
|
|
if err != nil {
|
|
t.Fatalf("invokeTool failed: %s", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != tc.wantCode {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
t.Errorf("status code got %d, want %d. Body: %s", resp.StatusCode, tc.wantCode, body)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func testError(t *testing.T, toolName string, request map[string]any, wantCode int, wantMsg string) {
|
|
resp, err := invokeTool(toolName, request, nil)
|
|
if err != nil {
|
|
t.Fatalf("invokeTool failed: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != wantCode {
|
|
bodyBytes, _ := io.ReadAll(resp.Body)
|
|
t.Fatalf("response status code is not %d, got %d: %s", wantCode, resp.StatusCode, string(bodyBytes))
|
|
}
|
|
|
|
bodyBytes, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
t.Fatalf("failed to read response body: %v", err)
|
|
}
|
|
|
|
if !bytes.Contains(bodyBytes, []byte(wantMsg)) {
|
|
t.Fatalf("response body does not contain %q: %s", wantMsg, string(bodyBytes))
|
|
}
|
|
}
|
|
|
|
func runListJobsTest(t *testing.T, client *dataproc.JobControllerClient, ctx context.Context) {
|
|
tcs := []struct {
|
|
name string
|
|
filter string
|
|
pageSize int
|
|
numPages int
|
|
wantN int
|
|
}{
|
|
{name: "one page", pageSize: 2, numPages: 1, wantN: 2},
|
|
{name: "two pages", pageSize: 1, numPages: 2, wantN: 2},
|
|
{name: "10 batches", pageSize: 10, numPages: 1, wantN: 10},
|
|
{name: "omit page size", numPages: 1, wantN: 20},
|
|
{
|
|
name: "filtered",
|
|
filter: "status.state = NON_ACTIVE",
|
|
pageSize: 20,
|
|
numPages: 1,
|
|
wantN: 20,
|
|
},
|
|
{
|
|
name: "empty",
|
|
filter: "status.state = NON_ACTIVE AND status.state = ACTIVE",
|
|
pageSize: 1,
|
|
numPages: 1,
|
|
wantN: 0,
|
|
},
|
|
}
|
|
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
want := []dataprocsrc.Job{}
|
|
if tc.wantN > 0 {
|
|
want = listJobsRpc(t, client, ctx, tc.filter, tc.wantN)
|
|
}
|
|
|
|
actual := []dataprocsrc.Job{}
|
|
var pageToken string
|
|
for i := 0; i < tc.numPages; i++ {
|
|
filter := tc.filter
|
|
if filter != "" {
|
|
filter += " AND "
|
|
}
|
|
filter += "clusterName = " + dataprocListJobsCluster
|
|
request := map[string]any{
|
|
"filter": filter,
|
|
"pageToken": pageToken,
|
|
}
|
|
if tc.pageSize > 0 {
|
|
request["pageSize"] = tc.pageSize
|
|
}
|
|
|
|
resp, err := invokeTool("list-jobs", request, nil)
|
|
if err != nil {
|
|
t.Fatalf("invokeTool failed: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
bodyBytes, _ := io.ReadAll(resp.Body)
|
|
t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
|
|
}
|
|
|
|
var body map[string]any
|
|
if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
|
|
t.Fatalf("error parsing response body: %v", err)
|
|
}
|
|
|
|
result, ok := body["result"].(string)
|
|
if !ok {
|
|
t.Fatalf("unable to find result in response body")
|
|
}
|
|
|
|
var listResponse dataprocsrc.ListJobsResponse
|
|
if err := json.Unmarshal([]byte(result), &listResponse); err != nil {
|
|
t.Fatalf("error unmarshalling result: %s", err)
|
|
}
|
|
actual = append(actual, listResponse.Jobs...)
|
|
pageToken = listResponse.NextPageToken
|
|
}
|
|
|
|
if !reflect.DeepEqual(actual, want) {
|
|
t.Fatalf("unexpected jobs: got %+v, want %+v", actual, want)
|
|
}
|
|
|
|
// want has URLs because it's created from Job instances by the same utility function
|
|
// used by the tool internals. Double-check that the URLs are reasonable.
|
|
for _, job := range want {
|
|
if !strings.HasPrefix(job.ConsoleURL, jobURLPrefix) {
|
|
t.Errorf("unexpected consoleUrl in job: %#v", job)
|
|
}
|
|
if !strings.HasPrefix(job.LogsURL, logsURLPrefix) {
|
|
t.Errorf("unexpected logsUrl in job: %#v", job)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func runGetJobTest(t *testing.T, client *dataproc.JobControllerClient, ctx context.Context, jobId string) {
|
|
// First get the job details directly from the Go proto API.
|
|
req := &dataprocpb.GetJobRequest{
|
|
ProjectId: dataprocProject,
|
|
Region: dataprocRegion,
|
|
JobId: jobId,
|
|
}
|
|
rawWantJobPb, err := client.GetJob(ctx, req)
|
|
if err != nil {
|
|
t.Fatalf("failed to get job: %s", err)
|
|
}
|
|
|
|
// Trim unknown fields from the proto by marshalling and unmarshalling.
|
|
jsonBytes, err := protojson.Marshal(rawWantJobPb)
|
|
if err != nil {
|
|
t.Fatalf("failed to marshal job to JSON: %s", err)
|
|
}
|
|
var wantJobPb dataprocpb.Job
|
|
if err := protojson.Unmarshal(jsonBytes, &wantJobPb); err != nil {
|
|
t.Fatalf("error unmarshalling result: %s", err)
|
|
}
|
|
|
|
tcs := []struct {
|
|
name string
|
|
jobId string
|
|
want *dataprocpb.Job
|
|
}{
|
|
{
|
|
name: "found job",
|
|
jobId: jobId,
|
|
want: &wantJobPb,
|
|
},
|
|
}
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
request := map[string]any{"jobId": tc.jobId}
|
|
resp, err := invokeTool("get-job", request, nil)
|
|
if err != nil {
|
|
t.Fatalf("invokeTool failed: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
t.Fatalf("status code got %d, want 200. Body: %s", resp.StatusCode, body)
|
|
}
|
|
|
|
var body map[string]any
|
|
if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
|
|
t.Fatalf("error parsing response body: %v", err)
|
|
}
|
|
resultStr, ok := body["result"].(string)
|
|
if !ok {
|
|
t.Fatalf("result is not a string, got %T", body["result"])
|
|
}
|
|
var wrappedResult map[string]any
|
|
if err := json.Unmarshal([]byte(resultStr), &wrappedResult); err != nil {
|
|
t.Fatalf("error unmarshalling result: %s", err)
|
|
}
|
|
|
|
consoleURL, ok := wrappedResult["consoleUrl"].(string)
|
|
if !ok || !strings.HasPrefix(consoleURL, jobURLPrefix) {
|
|
t.Errorf("unexpected consoleUrl: %v", consoleURL)
|
|
}
|
|
logsURL, ok := wrappedResult["logsUrl"].(string)
|
|
if !ok || !strings.HasPrefix(logsURL, logsURLPrefix) {
|
|
t.Errorf("unexpected logsUrl: %v", logsURL)
|
|
}
|
|
|
|
jobJSON, err := json.Marshal(wrappedResult["job"])
|
|
if err != nil {
|
|
t.Fatalf("failed to marshal job: %v", err)
|
|
}
|
|
|
|
// Unmarshal JSON to proto for proto-aware deep comparison.
|
|
var job dataprocpb.Job
|
|
if err := protojson.Unmarshal(jobJSON, &job); err != nil {
|
|
t.Fatalf("error unmarshalling job from wrapped result: %s", err)
|
|
}
|
|
|
|
if !cmp.Equal(&job, tc.want, protocmp.Transform()) {
|
|
diff := cmp.Diff(&job, tc.want, protocmp.Transform())
|
|
t.Errorf("GetJob() returned diff (-got +want):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("errors", func(t *testing.T) {
|
|
tcs := []struct {
|
|
name string
|
|
request map[string]any
|
|
wantCode int
|
|
wantMsg string
|
|
}{
|
|
{
|
|
name: "missing jobId",
|
|
request: map[string]any{},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: "missing required parameter: jobId",
|
|
},
|
|
{
|
|
name: "invalid name with slash",
|
|
request: map[string]any{"jobId": "projects/foo/regions/bar/jobs/baz"},
|
|
wantCode: http.StatusOK,
|
|
wantMsg: "jobId must be a short name without '/'",
|
|
},
|
|
}
|
|
for _, tc := range tcs {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
testError(t, "get-job", tc.request, tc.wantCode, tc.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
}
|
|
|
|
func listJobsRpc(t *testing.T, client *dataproc.JobControllerClient, ctx context.Context, filter string, n int) []dataprocsrc.Job {
|
|
req := &dataprocpb.ListJobsRequest{
|
|
ProjectId: dataprocProject,
|
|
Region: dataprocRegion,
|
|
ClusterName: dataprocListJobsCluster,
|
|
PageSize: int32(n),
|
|
}
|
|
if filter != "" {
|
|
req.Filter = filter
|
|
}
|
|
|
|
it := client.ListJobs(ctx, req)
|
|
pager := iterator.NewPager(it, n, "")
|
|
var jobPbs []*dataprocpb.Job
|
|
_, err := pager.NextPage(&jobPbs)
|
|
if err != nil {
|
|
t.Fatalf("failed to list jobs: %s", err)
|
|
}
|
|
|
|
jobs, err := dataprocsrc.ToJobs(jobPbs, dataprocRegion)
|
|
if err != nil {
|
|
t.Fatalf("failed to convert jobs: %v", err)
|
|
}
|
|
return jobs
|
|
}
|
|
|
|
func shortName(fullName string) string {
|
|
parts := strings.Split(fullName, "/")
|
|
return parts[len(parts)-1]
|
|
}
|