golang.org/x/build@v0.0.0-20240506185731-218518f32b70/internal/gomote/swarming_test.go (about)

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:build linux || darwin
     6  
     7  package gomote
     8  
     9  import (
    10  	"context"
    11  	"errors"
    12  	"fmt"
    13  	"io"
    14  	"log"
    15  	"net/http"
    16  	"net/http/httptest"
    17  	"os"
    18  	"strings"
    19  	"testing"
    20  	"time"
    21  
    22  	"github.com/google/go-cmp/cmp"
    23  	"github.com/google/uuid"
    24  	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
    25  	"go.chromium.org/luci/swarming/client/swarming"
    26  	"go.chromium.org/luci/swarming/client/swarming/swarmingtest"
    27  	swarmpb "go.chromium.org/luci/swarming/proto/api_v2"
    28  	"golang.org/x/build/internal/access"
    29  	"golang.org/x/build/internal/coordinator/remote"
    30  	"golang.org/x/build/internal/gomote/protos"
    31  	"golang.org/x/build/internal/rendezvous"
    32  	"golang.org/x/crypto/ssh"
    33  	"golang.org/x/net/nettest"
    34  	"google.golang.org/grpc"
    35  	"google.golang.org/grpc/codes"
    36  	"google.golang.org/grpc/status"
    37  	"google.golang.org/protobuf/testing/protocmp"
    38  )
    39  
    40  const testSwarmingBucketName = "unit-testing-bucket-swarming"
    41  
    42  func fakeGomoteSwarmingServer(t *testing.T, ctx context.Context, swarmClient swarming.Client, rdv rendezvousClient) protos.GomoteServiceServer {
    43  	signer, err := ssh.ParsePrivateKey([]byte(devCertCAPrivate))
    44  	if err != nil {
    45  		t.Fatalf("unable to parse raw certificate authority private key into signer=%s", err)
    46  	}
    47  	return &SwarmingServer{
    48  		bucket:                  &fakeBucketHandler{bucketName: testSwarmingBucketName},
    49  		buildlets:               remote.NewSessionPool(ctx),
    50  		gceBucketName:           testSwarmingBucketName,
    51  		sshCertificateAuthority: signer,
    52  		rendezvous:              rdv,
    53  		swarmingClient:          swarmClient,
    54  		buildersClient:          &FakeBuildersClient{},
    55  	}
    56  }
    57  
    58  func setupGomoteSwarmingTest(t *testing.T, ctx context.Context, swarmClient swarming.Client) protos.GomoteServiceClient {
    59  	lis, err := nettest.NewLocalListener("tcp")
    60  	if err != nil {
    61  		t.Fatalf("unable to create net listener: %s", err)
    62  	}
    63  	rdv := rendezvous.NewFake(context.Background(), func(ctx context.Context, jwt string) bool { return true })
    64  	sopts := access.FakeIAPAuthInterceptorOptions()
    65  	s := grpc.NewServer(sopts...)
    66  	protos.RegisterGomoteServiceServer(s, fakeGomoteSwarmingServer(t, ctx, swarmClient, rdv))
    67  	go s.Serve(lis)
    68  
    69  	// create GRPC client
    70  	copts := []grpc.DialOption{
    71  		grpc.WithInsecure(),
    72  		grpc.WithBlock(),
    73  		grpc.WithTimeout(5 * time.Second),
    74  	}
    75  	conn, err := grpc.Dial(lis.Addr().String(), copts...)
    76  	if err != nil {
    77  		lis.Close()
    78  		t.Fatalf("unable to create GRPC client: %s", err)
    79  	}
    80  	gc := protos.NewGomoteServiceClient(conn)
    81  	t.Cleanup(func() {
    82  		conn.Close()
    83  		s.Stop()
    84  		lis.Close()
    85  	})
    86  	return gc
    87  }
    88  
    89  func TestSwarmingAuthenticate(t *testing.T) {
    90  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
    91  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClient())
    92  	got, err := client.Authenticate(ctx, &protos.AuthenticateRequest{})
    93  	if err != nil {
    94  		t.Fatalf("client.Authenticate(ctx, request) = %v,  %s; want no error", got, err)
    95  	}
    96  }
    97  
    98  func TestSwarmingAuthenticateError(t *testing.T) {
    99  	wantCode := codes.Unauthenticated
   100  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClient())
   101  	_, err := client.Authenticate(context.Background(), &protos.AuthenticateRequest{})
   102  	if status.Code(err) != wantCode {
   103  		t.Fatalf("client.Authenticate(ctx, request) = _, %s; want %s", status.Code(err), wantCode)
   104  	}
   105  }
   106  
   107  func TestSwarmingAddBootstrap(t *testing.T) {
   108  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   109  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   110  	gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   111  	req := &protos.AddBootstrapRequest{
   112  		GomoteId: gomoteID,
   113  	}
   114  	got, err := client.AddBootstrap(ctx, req)
   115  	if err != nil {
   116  		t.Fatalf("client.AddBootstrap(ctx, %v) = %v, %s; want no error", req, got, err)
   117  	}
   118  }
   119  
   120  func TestSwarmingAddBootstrapError(t *testing.T) {
   121  	// This test will create a gomote instance and attempt to call AddBootstrap.
   122  	// If overrideID is set to true, the test will use a different gomoteID than
   123  	// the one created for the test.
   124  	testCases := []struct {
   125  		desc       string
   126  		ctx        context.Context
   127  		overrideID bool
   128  		gomoteID   string // Used iff overrideID is true.
   129  		wantCode   codes.Code
   130  	}{
   131  		{
   132  			desc:     "unauthenticated request",
   133  			ctx:      context.Background(),
   134  			wantCode: codes.Unauthenticated,
   135  		},
   136  		{
   137  			desc:       "missing gomote id",
   138  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   139  			overrideID: true,
   140  			wantCode:   codes.NotFound,
   141  		},
   142  		{
   143  			desc:       "gomote does not exist",
   144  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   145  			overrideID: true,
   146  			gomoteID:   "xyz",
   147  			wantCode:   codes.NotFound,
   148  		},
   149  		{
   150  			desc:     "gomote is not owned by caller",
   151  			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("user-x", "email-y")),
   152  			wantCode: codes.PermissionDenied,
   153  		},
   154  	}
   155  	for _, tc := range testCases {
   156  		t.Run(tc.desc, func(t *testing.T) {
   157  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   158  			gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   159  			if tc.overrideID {
   160  				gomoteID = tc.gomoteID
   161  			}
   162  			req := &protos.AddBootstrapRequest{
   163  				GomoteId: gomoteID,
   164  			}
   165  			got, err := client.AddBootstrap(tc.ctx, req)
   166  			if err != nil && status.Code(err) != tc.wantCode {
   167  				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
   168  			}
   169  			if err == nil {
   170  				t.Fatalf("client.AddBootstrap(ctx, %v) = %v, nil; want error", req, got)
   171  			}
   172  		})
   173  	}
   174  }
   175  
   176  func TestSwarmingListSwarmingBuilders(t *testing.T) {
   177  	log.SetOutput(io.Discard)
   178  	defer log.SetOutput(os.Stdout)
   179  
   180  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClient())
   181  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   182  	response, err := client.ListSwarmingBuilders(ctx, &protos.ListSwarmingBuildersRequest{})
   183  	if err != nil {
   184  		t.Fatalf("client.ListSwarmingBuilders = nil, %s; want no error", err)
   185  	}
   186  	got := response.GetBuilders()
   187  	if diff := cmp.Diff([]string{"gotip-linux-amd64", "gotip-linux-amd64-boringcrypto", "gotip-linux-arm"}, got); diff != "" {
   188  		t.Errorf("ListBuilders() mismatch (-want, +got):\n%s", diff)
   189  	}
   190  }
   191  
   192  func TestSwarmingListSwarmingBuildersError(t *testing.T) {
   193  	log.SetOutput(io.Discard)
   194  	defer log.SetOutput(os.Stdout)
   195  
   196  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClient())
   197  	req := &protos.ListSwarmingBuildersRequest{}
   198  	got, err := client.ListSwarmingBuilders(context.Background(), req)
   199  	if err != nil && status.Code(err) != codes.Unauthenticated {
   200  		t.Fatalf("unexpected error: %s; want %s", err, codes.Unauthenticated)
   201  	}
   202  	if err == nil {
   203  		t.Fatalf("client.ListSwarmingBuilder(ctx, %v) = %v, nil; want error", req, got)
   204  	}
   205  }
   206  
   207  func TestSwarmingCreateInstance(t *testing.T) {
   208  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   209  	req := &protos.CreateInstanceRequest{BuilderType: "gotip-linux-amd64-boringcrypto"}
   210  
   211  	msc := mockSwarmClient()
   212  	msc.NewTaskMock = func(_ context.Context, req *swarmpb.NewTaskRequest) (*swarmpb.TaskRequestMetadataResponse, error) {
   213  		taskID := uuid.New().String()
   214  		return &swarmpb.TaskRequestMetadataResponse{
   215  			TaskId: taskID,
   216  			Request: &swarmpb.TaskRequestResponse{
   217  				TaskId: taskID,
   218  				Name:   req.Name,
   219  			},
   220  		}, nil
   221  	}
   222  	msc.TaskResultMock = func(_ context.Context, taskID string, _ *swarming.TaskResultFields) (*swarmpb.TaskResultResponse, error) {
   223  		return &swarmpb.TaskResultResponse{
   224  			TaskId: taskID,
   225  			State:  swarmpb.TaskState_RUNNING,
   226  		}, nil
   227  	}
   228  
   229  	gomoteClient := setupGomoteSwarmingTest(t, context.Background(), msc)
   230  
   231  	stream, err := gomoteClient.CreateInstance(ctx, req)
   232  	if err != nil {
   233  		t.Fatalf("client.CreateInstance(ctx, %v) = %v,  %s; want no error", req, stream, err)
   234  	}
   235  	var updateComplete bool
   236  	for {
   237  		update, err := stream.Recv()
   238  		if err == io.EOF && !updateComplete {
   239  			t.Fatal("stream.Recv = stream, io.EOF; want no EOF")
   240  		}
   241  		if err == io.EOF {
   242  			break
   243  		}
   244  		if err != nil {
   245  			t.Fatalf("stream.Recv() = nil, %s; want no error", err)
   246  		}
   247  		if update.GetStatus() == protos.CreateInstanceResponse_COMPLETE {
   248  			updateComplete = true
   249  		}
   250  	}
   251  }
   252  
   253  func TestSwarmingCreateInstanceError(t *testing.T) {
   254  	log.SetOutput(io.Discard)
   255  	defer log.SetOutput(os.Stdout)
   256  
   257  	testCases := []struct {
   258  		desc     string
   259  		ctx      context.Context
   260  		request  *protos.CreateInstanceRequest
   261  		wantCode codes.Code
   262  	}{
   263  		{
   264  			desc:     "unauthenticated request",
   265  			ctx:      context.Background(),
   266  			request:  &protos.CreateInstanceRequest{},
   267  			wantCode: codes.Unauthenticated,
   268  		},
   269  		{
   270  			desc:     "missing builder type",
   271  			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   272  			request:  &protos.CreateInstanceRequest{},
   273  			wantCode: codes.InvalidArgument,
   274  		},
   275  		{
   276  			desc: "invalid builder type",
   277  			ctx:  access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   278  			request: &protos.CreateInstanceRequest{
   279  				BuilderType: "funky-time-builder",
   280  			},
   281  			wantCode: codes.InvalidArgument,
   282  		},
   283  	}
   284  	for _, tc := range testCases {
   285  		t.Run(tc.desc, func(t *testing.T) {
   286  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClient())
   287  
   288  			stream, err := client.CreateInstance(tc.ctx, tc.request)
   289  			if err != nil {
   290  				t.Fatalf("client.CreateInstance(ctx, %v) = %v,  %s; want no error", tc.request, stream, err)
   291  			}
   292  			for {
   293  				_, got := stream.Recv()
   294  				if got == io.EOF {
   295  					t.Fatal("stream.Recv = stream, io.EOF; want no EOF")
   296  				}
   297  				if got != nil && status.Code(got) != tc.wantCode {
   298  					t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
   299  				}
   300  				return
   301  			}
   302  		})
   303  	}
   304  }
   305  
   306  func TestSwarmingDestroyInstance(t *testing.T) {
   307  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   308  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   309  	gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   310  	if _, err := client.DestroyInstance(ctx, &protos.DestroyInstanceRequest{
   311  		GomoteId: gomoteID,
   312  	}); err != nil {
   313  		t.Fatalf("client.DestroyInstance(ctx, req) = response, %s; want no error", err)
   314  	}
   315  }
   316  
   317  func TestSwarmingDestroyInstanceError(t *testing.T) {
   318  	// This test will create a gomote instance and attempt to call DestroyInstance.
   319  	// If overrideID is set to true, the test will use a different gomoteID than
   320  	// the one created for the test.
   321  	testCases := []struct {
   322  		desc       string
   323  		ctx        context.Context
   324  		overrideID bool
   325  		gomoteID   string // Used iff overrideID is true.
   326  		wantCode   codes.Code
   327  	}{
   328  		{
   329  			desc:     "unauthenticated request",
   330  			ctx:      context.Background(),
   331  			wantCode: codes.Unauthenticated,
   332  		},
   333  		{
   334  			desc:       "missing gomote id",
   335  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   336  			overrideID: true,
   337  			gomoteID:   "",
   338  			wantCode:   codes.InvalidArgument,
   339  		},
   340  		{
   341  			desc:       "gomote does not exist",
   342  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   343  			overrideID: true,
   344  			gomoteID:   "chucky",
   345  			wantCode:   codes.NotFound,
   346  		},
   347  		{
   348  			desc:       "wrong gomote id",
   349  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   350  			overrideID: false,
   351  			wantCode:   codes.PermissionDenied,
   352  		},
   353  	}
   354  	for _, tc := range testCases {
   355  		t.Run(tc.desc, func(t *testing.T) {
   356  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   357  			gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   358  			if tc.overrideID {
   359  				gomoteID = tc.gomoteID
   360  			}
   361  			req := &protos.DestroyInstanceRequest{
   362  				GomoteId: gomoteID,
   363  			}
   364  			got, err := client.DestroyInstance(tc.ctx, req)
   365  			if err != nil && status.Code(err) != tc.wantCode {
   366  				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
   367  			}
   368  			if err == nil {
   369  				t.Fatalf("client.DestroyInstance(ctx, %v) = %v, nil; want error", req, got)
   370  			}
   371  		})
   372  	}
   373  }
   374  
   375  func TestSwarmingExecuteCommand(t *testing.T) {
   376  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   377  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   378  	gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   379  	stream, err := client.ExecuteCommand(ctx, &protos.ExecuteCommandRequest{
   380  		GomoteId:          gomoteID,
   381  		Command:           "ls",
   382  		SystemLevel:       false,
   383  		Debug:             false,
   384  		AppendEnvironment: nil,
   385  		Path:              nil,
   386  		Directory:         "/workdir",
   387  		Args:              []string{"-alh"},
   388  	})
   389  	if err != nil {
   390  		t.Fatalf("client.ExecuteCommand(ctx, req) = response, %s; want no error", err)
   391  	}
   392  	var out []byte
   393  	for {
   394  		res, err := stream.Recv()
   395  		if err != nil && err == io.EOF {
   396  			break
   397  		}
   398  		if err != nil {
   399  			t.Fatalf("stream.Recv() = _, %s; want no error", err)
   400  		}
   401  		out = append(out, res.GetOutput()...)
   402  	}
   403  	if len(out) == 0 {
   404  		t.Fatalf("output: %q, expected non-empty", out)
   405  	}
   406  }
   407  
   408  func TestSwarmingExecuteCommandError(t *testing.T) {
   409  	// This test will create a gomote instance and attempt to call TestExecuteCommand.
   410  	// If overrideID is set to true, the test will use a different gomoteID than
   411  	// the one created for the test.
   412  	testCases := []struct {
   413  		desc       string
   414  		ctx        context.Context
   415  		overrideID bool
   416  		gomoteID   string // Used iff overrideID is true.
   417  		cmd        string
   418  		wantCode   codes.Code
   419  	}{
   420  		{
   421  			desc:     "unauthenticated request",
   422  			ctx:      context.Background(),
   423  			wantCode: codes.Unauthenticated,
   424  		},
   425  		{
   426  			desc:       "missing gomote id",
   427  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   428  			overrideID: true,
   429  			gomoteID:   "",
   430  			wantCode:   codes.NotFound,
   431  		},
   432  		{
   433  			desc:     "missing command",
   434  			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   435  			wantCode: codes.Aborted,
   436  		},
   437  		{
   438  			desc:       "gomote does not exist",
   439  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   440  			overrideID: true,
   441  			gomoteID:   "chucky",
   442  			cmd:        "ls",
   443  			wantCode:   codes.NotFound,
   444  		},
   445  		{
   446  			desc:       "wrong gomote id",
   447  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   448  			overrideID: false,
   449  			cmd:        "ls",
   450  			wantCode:   codes.PermissionDenied,
   451  		},
   452  	}
   453  	for _, tc := range testCases {
   454  		t.Run(tc.desc, func(t *testing.T) {
   455  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   456  			gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   457  			if tc.overrideID {
   458  				gomoteID = tc.gomoteID
   459  			}
   460  			stream, err := client.ExecuteCommand(tc.ctx, &protos.ExecuteCommandRequest{
   461  				GomoteId:          gomoteID,
   462  				Command:           tc.cmd,
   463  				SystemLevel:       false,
   464  				Debug:             false,
   465  				AppendEnvironment: nil,
   466  				Path:              nil,
   467  				Directory:         "/workdir",
   468  				Args:              []string{"-alh"},
   469  			})
   470  			if err != nil {
   471  				t.Fatalf("unexpected error: %s", err)
   472  			}
   473  			res, err := stream.Recv()
   474  			if err != nil && status.Code(err) != tc.wantCode {
   475  				t.Fatalf("unexpected error: %s", err)
   476  			}
   477  			if err == nil {
   478  				t.Fatalf("client.ExecuteCommand(ctx, req) = %v, nil; want error", res)
   479  			}
   480  		})
   481  	}
   482  }
   483  
   484  func TestSwarmingInstanceAlive(t *testing.T) {
   485  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   486  	gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   487  	req := &protos.InstanceAliveRequest{
   488  		GomoteId: gomoteID,
   489  	}
   490  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   491  	got, err := client.InstanceAlive(ctx, req)
   492  	if err != nil {
   493  		t.Fatalf("client.InstanceAlive(ctx, %v) = %v, %s; want no error", req, got, err)
   494  	}
   495  }
   496  
   497  func TestSwarmingInstanceAliveError(t *testing.T) {
   498  	// This test will create a gomote instance and attempt to call InstanceAlive.
   499  	// If overrideID is set to true, the test will use a different gomoteID than
   500  	// the one created for the test.
   501  	testCases := []struct {
   502  		desc       string
   503  		ctx        context.Context
   504  		overrideID bool
   505  		gomoteID   string // Used iff overrideID is true.
   506  		wantCode   codes.Code
   507  	}{
   508  		{
   509  			desc:     "unauthenticated request",
   510  			ctx:      context.Background(),
   511  			wantCode: codes.Unauthenticated,
   512  		},
   513  		{
   514  			desc:       "missing gomote id",
   515  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   516  			overrideID: true,
   517  			wantCode:   codes.InvalidArgument,
   518  		},
   519  		{
   520  			desc:       "gomote does not exist",
   521  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   522  			overrideID: true,
   523  			gomoteID:   "xyz",
   524  			wantCode:   codes.NotFound,
   525  		},
   526  		{
   527  			desc:     "gomote is not owned by caller",
   528  			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("user-x", "email-y")),
   529  			wantCode: codes.PermissionDenied,
   530  		},
   531  	}
   532  	for _, tc := range testCases {
   533  		t.Run(tc.desc, func(t *testing.T) {
   534  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   535  			gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   536  			if tc.overrideID {
   537  				gomoteID = tc.gomoteID
   538  			}
   539  			req := &protos.InstanceAliveRequest{
   540  				GomoteId: gomoteID,
   541  			}
   542  			got, err := client.InstanceAlive(tc.ctx, req)
   543  			if err != nil && status.Code(err) != tc.wantCode {
   544  				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
   545  			}
   546  			if err == nil {
   547  				t.Fatalf("client.InstanceAlive(ctx, %v) = %v, nil; want error", req, got)
   548  			}
   549  		})
   550  	}
   551  }
   552  
   553  func TestSwarmingListDirectory(t *testing.T) {
   554  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   555  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   556  	gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   557  	if _, err := client.ListDirectory(ctx, &protos.ListDirectoryRequest{
   558  		GomoteId:  gomoteID,
   559  		Directory: "/foo",
   560  	}); err != nil {
   561  		t.Fatalf("client.RemoveFiles(ctx, req) = response, %s; want no error", err)
   562  	}
   563  }
   564  
   565  func TestSwarmingListDirectoryError(t *testing.T) {
   566  	// This test will create a gomote instance and attempt to call ListDirectory.
   567  	// If overrideID is set to true, the test will use a different gomoteID than
   568  	// the one created for the test.
   569  	testCases := []struct {
   570  		desc       string
   571  		ctx        context.Context
   572  		overrideID bool
   573  		gomoteID   string // Used iff overrideID is true.
   574  		directory  string
   575  		recursive  bool
   576  		skipFiles  []string
   577  		digest     bool
   578  		wantCode   codes.Code
   579  	}{
   580  		{
   581  			desc:     "unauthenticated request",
   582  			ctx:      context.Background(),
   583  			wantCode: codes.Unauthenticated,
   584  		},
   585  		{
   586  			desc:       "missing gomote id",
   587  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   588  			overrideID: true,
   589  			gomoteID:   "",
   590  			wantCode:   codes.InvalidArgument,
   591  		},
   592  		{
   593  			desc:     "missing directory",
   594  			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   595  			wantCode: codes.InvalidArgument,
   596  		},
   597  		{
   598  			desc:       "gomote does not exist",
   599  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   600  			overrideID: true,
   601  			gomoteID:   "chucky",
   602  			directory:  "/foo",
   603  			wantCode:   codes.NotFound,
   604  		},
   605  		{
   606  			desc:       "wrong gomote id",
   607  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   608  			overrideID: false,
   609  			directory:  "/foo",
   610  			wantCode:   codes.PermissionDenied,
   611  		},
   612  	}
   613  	for _, tc := range testCases {
   614  		t.Run(tc.desc, func(t *testing.T) {
   615  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   616  			gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   617  			if tc.overrideID {
   618  				gomoteID = tc.gomoteID
   619  			}
   620  			req := &protos.ListDirectoryRequest{
   621  				GomoteId:  gomoteID,
   622  				Directory: tc.directory,
   623  				Recursive: false,
   624  				SkipFiles: []string{},
   625  				Digest:    false,
   626  			}
   627  			got, err := client.ListDirectory(tc.ctx, req)
   628  			if err != nil && status.Code(err) != tc.wantCode {
   629  				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
   630  			}
   631  			if err == nil {
   632  				t.Fatalf("client.RemoveFiles(ctx, %v) = %v, nil; want error", req, got)
   633  			}
   634  		})
   635  	}
   636  }
   637  
   638  func TestSwarmingListInstance(t *testing.T) {
   639  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   640  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   641  	var want []*protos.Instance
   642  	for i := 0; i < 3; i++ {
   643  		want = append(want, &protos.Instance{
   644  			GomoteId:    mustCreateSwarmingInstance(t, client, fakeIAP()),
   645  			BuilderType: "gotip-linux-amd64-boringcrypto",
   646  		})
   647  	}
   648  	mustCreateSwarmingInstance(t, client, fakeIAPWithUser("user-x", "uuid-user-x"))
   649  	response, err := client.ListInstances(ctx, &protos.ListInstancesRequest{})
   650  	if err != nil {
   651  		t.Fatalf("client.ListInstances = nil, %s; want no error", err)
   652  	}
   653  	got := response.GetInstances()
   654  	if diff := cmp.Diff(want, got, protocmp.Transform(), protocmp.IgnoreFields(&protos.Instance{}, "expires", "host_type")); diff != "" {
   655  		t.Errorf("ListInstances() mismatch (-want, +got):\n%s", diff)
   656  	}
   657  }
   658  
   659  func TestSwarmingReadTGZToURLError(t *testing.T) {
   660  	// This test will create a gomote instance and attempt to call ReadTGZToURL.
   661  	// If overrideID is set to true, the test will use a different gomoteID than
   662  	// the one created for the test.
   663  	testCases := []struct {
   664  		desc       string
   665  		ctx        context.Context
   666  		overrideID bool
   667  		gomoteID   string // Used iff overrideID is true.
   668  		directory  string
   669  		wantCode   codes.Code
   670  	}{
   671  		{
   672  			desc:     "unauthenticated request",
   673  			ctx:      context.Background(),
   674  			wantCode: codes.Unauthenticated,
   675  		},
   676  		{
   677  			desc:       "missing gomote id",
   678  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   679  			overrideID: true,
   680  			gomoteID:   "",
   681  			wantCode:   codes.NotFound,
   682  		},
   683  		{
   684  			desc:       "gomote does not exist",
   685  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   686  			overrideID: true,
   687  			gomoteID:   "chucky",
   688  			wantCode:   codes.NotFound,
   689  		},
   690  		{
   691  			desc:       "wrong gomote id",
   692  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   693  			overrideID: false,
   694  			wantCode:   codes.PermissionDenied,
   695  		},
   696  	}
   697  	for _, tc := range testCases {
   698  		t.Run(tc.desc, func(t *testing.T) {
   699  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   700  			gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   701  			if tc.overrideID {
   702  				gomoteID = tc.gomoteID
   703  			}
   704  			req := &protos.ReadTGZToURLRequest{
   705  				GomoteId:  gomoteID,
   706  				Directory: tc.directory,
   707  			}
   708  			got, err := client.ReadTGZToURL(tc.ctx, req)
   709  			if err != nil && status.Code(err) != tc.wantCode {
   710  				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
   711  			}
   712  			if err == nil {
   713  				t.Fatalf("client.ReadTGZToURL(ctx, %v) = %v, nil; want error", req, got)
   714  			}
   715  		})
   716  	}
   717  }
   718  
   719  func TestSwarmingRemoveFiles(t *testing.T) {
   720  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   721  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   722  	gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   723  	if _, err := client.RemoveFiles(ctx, &protos.RemoveFilesRequest{
   724  		GomoteId: gomoteID,
   725  		Paths:    []string{"temp_file.log"},
   726  	}); err != nil {
   727  		t.Fatalf("client.RemoveFiles(ctx, req) = response, %s; want no error", err)
   728  	}
   729  }
   730  
   731  func TestSwarmingSignSSHKey(t *testing.T) {
   732  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   733  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   734  	gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   735  	if _, err := client.SignSSHKey(ctx, &protos.SignSSHKeyRequest{
   736  		GomoteId:     gomoteID,
   737  		PublicSshKey: []byte(devCertCAPublic),
   738  	}); err != nil {
   739  		t.Fatalf("client.SignSSHKey(ctx, req) = response, %s; want no error", err)
   740  	}
   741  }
   742  
   743  func TestSwarmingSignSSHKeyError(t *testing.T) {
   744  	// This test will create a gomote instance and attempt to call SignSSHKey.
   745  	// If overrideID is set to true, the test will use a different gomoteID than
   746  	// the one created for the test.
   747  	testCases := []struct {
   748  		desc          string
   749  		ctx           context.Context
   750  		overrideID    bool
   751  		gomoteID      string // Used iff overrideID is true.
   752  		publickSSHKey []byte
   753  		wantCode      codes.Code
   754  	}{
   755  		{
   756  			desc:     "unauthenticated request",
   757  			ctx:      context.Background(),
   758  			wantCode: codes.Unauthenticated,
   759  		},
   760  		{
   761  			desc:       "missing gomote id",
   762  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   763  			overrideID: true,
   764  			gomoteID:   "",
   765  			wantCode:   codes.NotFound,
   766  		},
   767  		{
   768  			desc:     "missing public key",
   769  			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   770  			wantCode: codes.InvalidArgument,
   771  		},
   772  		{
   773  			desc:          "gomote does not exist",
   774  			ctx:           access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   775  			overrideID:    true,
   776  			gomoteID:      "chucky",
   777  			publickSSHKey: []byte(devCertCAPublic),
   778  			wantCode:      codes.NotFound,
   779  		},
   780  		{
   781  			desc:          "wrong gomote id",
   782  			ctx:           access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   783  			overrideID:    false,
   784  			publickSSHKey: []byte(devCertCAPublic),
   785  			wantCode:      codes.PermissionDenied,
   786  		},
   787  	}
   788  	for _, tc := range testCases {
   789  		t.Run(tc.desc, func(t *testing.T) {
   790  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   791  			gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   792  			if tc.overrideID {
   793  				gomoteID = tc.gomoteID
   794  			}
   795  			req := &protos.SignSSHKeyRequest{
   796  				GomoteId:     gomoteID,
   797  				PublicSshKey: tc.publickSSHKey,
   798  			}
   799  			got, err := client.SignSSHKey(tc.ctx, req)
   800  			if err != nil && status.Code(err) != tc.wantCode {
   801  				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
   802  			}
   803  			if err == nil {
   804  				t.Fatalf("client.SignSSHKey(ctx, %v) = %v, nil; want error", req, got)
   805  			}
   806  		})
   807  	}
   808  }
   809  
   810  func TestSwarmingUploadFile(t *testing.T) {
   811  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   812  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   813  	_ = mustCreateSwarmingInstance(t, client, fakeIAP())
   814  	if _, err := client.UploadFile(ctx, &protos.UploadFileRequest{}); err != nil {
   815  		t.Fatalf("client.UploadFile(ctx, req) = response, %s; want no error", err)
   816  	}
   817  }
   818  
   819  func TestSwarmingUploadFileError(t *testing.T) {
   820  	// This test will create a gomote instance and attempt to call UploadFile.
   821  	// If overrideID is set to true, the test will use a different gomoteID than
   822  	// the one created for the test.
   823  	testCases := []struct {
   824  		desc       string
   825  		ctx        context.Context
   826  		overrideID bool
   827  		filename   string
   828  		wantCode   codes.Code
   829  	}{
   830  		{
   831  			desc:     "unauthenticated request",
   832  			ctx:      context.Background(),
   833  			wantCode: codes.Unauthenticated,
   834  		},
   835  	}
   836  	for _, tc := range testCases {
   837  		t.Run(tc.desc, func(t *testing.T) {
   838  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   839  			_ = mustCreateSwarmingInstance(t, client, fakeIAP())
   840  			req := &protos.UploadFileRequest{}
   841  			got, err := client.UploadFile(tc.ctx, req)
   842  			if err != nil && status.Code(err) != tc.wantCode {
   843  				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
   844  			}
   845  			if err == nil {
   846  				t.Fatalf("client.UploadFile(ctx, %v) = %v, nil; want error", req, got)
   847  			}
   848  		})
   849  	}
   850  }
   851  
   852  func TestSwarmingRemoveFilesError(t *testing.T) {
   853  	// This test will create a gomote instance and attempt to call RemoveFiles.
   854  	// If overrideID is set to true, the test will use a different gomoteID than
   855  	// the one created for the test.
   856  	testCases := []struct {
   857  		desc       string
   858  		ctx        context.Context
   859  		overrideID bool
   860  		gomoteID   string // Used iff overrideID is true.
   861  		paths      []string
   862  		wantCode   codes.Code
   863  	}{
   864  		{
   865  			desc:     "unauthenticated request",
   866  			ctx:      context.Background(),
   867  			wantCode: codes.Unauthenticated,
   868  		},
   869  		{
   870  			desc:       "missing gomote id",
   871  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   872  			overrideID: true,
   873  			gomoteID:   "",
   874  			wantCode:   codes.InvalidArgument,
   875  		},
   876  		{
   877  			desc:     "missing paths",
   878  			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   879  			paths:    []string{},
   880  			wantCode: codes.InvalidArgument,
   881  		},
   882  		{
   883  			desc:       "gomote does not exist",
   884  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   885  			overrideID: true,
   886  			gomoteID:   "chucky",
   887  			paths:      []string{"file.a"},
   888  			wantCode:   codes.NotFound,
   889  		},
   890  		{
   891  			desc:       "wrong gomote id",
   892  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   893  			overrideID: false,
   894  			paths:      []string{"file.a"},
   895  			wantCode:   codes.PermissionDenied,
   896  		},
   897  	}
   898  	for _, tc := range testCases {
   899  		t.Run(tc.desc, func(t *testing.T) {
   900  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   901  			gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   902  			if tc.overrideID {
   903  				gomoteID = tc.gomoteID
   904  			}
   905  			req := &protos.RemoveFilesRequest{
   906  				GomoteId: gomoteID,
   907  				Paths:    tc.paths,
   908  			}
   909  			got, err := client.RemoveFiles(tc.ctx, req)
   910  			if err != nil && status.Code(err) != tc.wantCode {
   911  				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
   912  			}
   913  			if err == nil {
   914  				t.Fatalf("client.RemoveFiles(ctx, %v) = %v, nil; want error", req, got)
   915  			}
   916  		})
   917  	}
   918  }
   919  
   920  // TODO(go.dev/issue/48737) add test for files on GCS
   921  func TestSwarmingWriteFileFromURL(t *testing.T) {
   922  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
   923  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   924  	gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   925  	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   926  		fmt.Fprintln(w, "Go is an open source programming language")
   927  	}))
   928  	defer ts.Close()
   929  	if _, err := client.WriteFileFromURL(ctx, &protos.WriteFileFromURLRequest{
   930  		GomoteId: gomoteID,
   931  		Url:      ts.URL,
   932  		Filename: "foo",
   933  		Mode:     0777,
   934  	}); err != nil {
   935  		t.Fatalf("client.WriteFileFromURL(ctx, req) = response, %s; want no error", err)
   936  	}
   937  }
   938  
   939  func TestSwarmingWriteFileFromURLError(t *testing.T) {
   940  	// This test will create a gomote instance and attempt to call TestWriteFileFromURL.
   941  	// If overrideID is set to true, the test will use a different gomoteID than
   942  	// the one created for the test.
   943  	testCases := []struct {
   944  		desc       string
   945  		ctx        context.Context
   946  		overrideID bool
   947  		gomoteID   string // Used iff overrideID is true.
   948  		url        string
   949  		filename   string
   950  		mode       uint32
   951  		wantCode   codes.Code
   952  	}{
   953  		{
   954  			desc:     "unauthenticated request",
   955  			ctx:      context.Background(),
   956  			wantCode: codes.Unauthenticated,
   957  		},
   958  		{
   959  			desc:       "missing gomote id",
   960  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
   961  			overrideID: true,
   962  			gomoteID:   "",
   963  			wantCode:   codes.NotFound,
   964  		},
   965  		{
   966  			desc:       "gomote does not exist",
   967  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   968  			overrideID: true,
   969  			gomoteID:   "chucky",
   970  			url:        "go.dev/dl/1_14.tar.gz",
   971  			wantCode:   codes.NotFound,
   972  		},
   973  		{
   974  			desc:       "wrong gomote id",
   975  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
   976  			overrideID: false,
   977  			url:        "go.dev/dl/1_14.tar.gz",
   978  			wantCode:   codes.PermissionDenied,
   979  		},
   980  	}
   981  	for _, tc := range testCases {
   982  		t.Run(tc.desc, func(t *testing.T) {
   983  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
   984  			gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
   985  			if tc.overrideID {
   986  				gomoteID = tc.gomoteID
   987  			}
   988  			req := &protos.WriteFileFromURLRequest{
   989  				GomoteId: gomoteID,
   990  				Url:      tc.url,
   991  				Filename: tc.filename,
   992  				Mode:     0,
   993  			}
   994  			got, err := client.WriteFileFromURL(tc.ctx, req)
   995  			if err != nil && status.Code(err) != tc.wantCode {
   996  				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
   997  			}
   998  			if err == nil {
   999  				t.Fatalf("client.WriteFileFromURL(ctx, %v) = %v, nil; want error", req, got)
  1000  			}
  1001  		})
  1002  	}
  1003  }
  1004  
  1005  func TestSwarmingWriteTGZFromURL(t *testing.T) {
  1006  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
  1007  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
  1008  	gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
  1009  	if _, err := client.WriteTGZFromURL(ctx, &protos.WriteTGZFromURLRequest{
  1010  		GomoteId:  gomoteID,
  1011  		Directory: "foo",
  1012  		Url:       `https://go.dev/dl/go1.17.6.linux-amd64.tar.gz`,
  1013  	}); err != nil {
  1014  		t.Fatalf("client.WriteTGZFromURL(ctx, req) = response, %s; want no error", err)
  1015  	}
  1016  }
  1017  
  1018  func TestSwarmingWriteTGZFromURLGomoteStaging(t *testing.T) {
  1019  	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
  1020  	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
  1021  	gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
  1022  	if _, err := client.WriteTGZFromURL(ctx, &protos.WriteTGZFromURLRequest{
  1023  		GomoteId:  gomoteID,
  1024  		Directory: "foo",
  1025  		Url:       fmt.Sprintf("https://storage.googleapis.com/%s/go1.17.6.linux-amd64.tar.gz?field=x", testBucketName),
  1026  	}); err != nil {
  1027  		t.Fatalf("client.WriteTGZFromURL(ctx, req) = response, %s; want no error", err)
  1028  	}
  1029  }
  1030  
  1031  func TestSwarmingWriteTGZFromURLError(t *testing.T) {
  1032  	// This test will create a gomote instance and attempt to call TestWriteTGZFromURL.
  1033  	// If overrideID is set to true, the test will use a different gomoteID than
  1034  	// the one created for the test.
  1035  	testCases := []struct {
  1036  		desc       string
  1037  		ctx        context.Context
  1038  		overrideID bool
  1039  		gomoteID   string // Used iff overrideID is true.
  1040  		url        string
  1041  		directory  string
  1042  		wantCode   codes.Code
  1043  	}{
  1044  		{
  1045  			desc:     "unauthenticated request",
  1046  			ctx:      context.Background(),
  1047  			wantCode: codes.Unauthenticated,
  1048  		},
  1049  		{
  1050  			desc:       "missing gomote id",
  1051  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
  1052  			overrideID: true,
  1053  			gomoteID:   "",
  1054  			wantCode:   codes.InvalidArgument,
  1055  		},
  1056  		{
  1057  			desc:     "missing URL",
  1058  			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
  1059  			wantCode: codes.InvalidArgument,
  1060  		},
  1061  		{
  1062  			desc:       "gomote does not exist",
  1063  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
  1064  			overrideID: true,
  1065  			gomoteID:   "chucky",
  1066  			url:        "go.dev/dl/1_14.tar.gz",
  1067  			wantCode:   codes.NotFound,
  1068  		},
  1069  		{
  1070  			desc:       "wrong gomote id",
  1071  			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
  1072  			overrideID: false,
  1073  			url:        "go.dev/dl/1_14.tar.gz",
  1074  			wantCode:   codes.PermissionDenied,
  1075  		},
  1076  	}
  1077  	for _, tc := range testCases {
  1078  		t.Run(tc.desc, func(t *testing.T) {
  1079  			client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
  1080  			gomoteID := mustCreateSwarmingInstance(t, client, fakeIAP())
  1081  			if tc.overrideID {
  1082  				gomoteID = tc.gomoteID
  1083  			}
  1084  			req := &protos.WriteTGZFromURLRequest{
  1085  				GomoteId:  gomoteID,
  1086  				Url:       tc.url,
  1087  				Directory: tc.directory,
  1088  			}
  1089  			got, err := client.WriteTGZFromURL(tc.ctx, req)
  1090  			if err != nil && status.Code(err) != tc.wantCode {
  1091  				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
  1092  			}
  1093  			if err == nil {
  1094  				t.Fatalf("client.WriteTGZFromURL(ctx, %v) = %v, nil; want error", req, got)
  1095  			}
  1096  		})
  1097  	}
  1098  }
  1099  
  1100  func TestStartNewSwarmingTask(t *testing.T) {
  1101  	log.SetOutput(io.Discard)
  1102  	defer log.SetOutput(os.Stdout)
  1103  
  1104  	ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
  1105  	rdv := rendezvous.New(ctx, rendezvous.OptionValidator(func(ctx context.Context, jwt string) bool {
  1106  		return true
  1107  	}))
  1108  	ts := httptest.NewTLSServer(http.HandlerFunc(rdv.HandleReverse))
  1109  	defer ts.Close()
  1110  
  1111  	msc := mockSwarmClient()
  1112  	msc.NewTaskMock = func(_ context.Context, req *swarmpb.NewTaskRequest) (*swarmpb.TaskRequestMetadataResponse, error) {
  1113  		taskID := uuid.New().String()
  1114  		return &swarmpb.TaskRequestMetadataResponse{
  1115  			TaskId: taskID,
  1116  			Request: &swarmpb.TaskRequestResponse{
  1117  				TaskId: taskID,
  1118  				Name:   req.Name,
  1119  			},
  1120  		}, nil
  1121  	}
  1122  	msc.TaskResultMock = func(_ context.Context, taskID string, _ *swarming.TaskResultFields) (*swarmpb.TaskResultResponse, error) {
  1123  		return &swarmpb.TaskResultResponse{
  1124  			TaskId: taskID,
  1125  			State:  swarmpb.TaskState_RUNNING,
  1126  		}, nil
  1127  	}
  1128  	ss := &SwarmingServer{
  1129  		UnimplementedGomoteServiceServer: protos.UnimplementedGomoteServiceServer{},
  1130  		bucket:                           nil,
  1131  		buildlets:                        &remote.SessionPool{},
  1132  		gceBucketName:                    "",
  1133  		sshCertificateAuthority:          nil,
  1134  		rendezvous:                       rdv,
  1135  		swarmingClient:                   msc,
  1136  		buildersClient:                   &FakeBuildersClient{},
  1137  	}
  1138  	id := "task-123"
  1139  	errCh := make(chan error, 2)
  1140  	if _, err := ss.startNewSwarmingTask(ctx, id, map[string]string{"cipd_platform": "linux-amd64"}, &configProperties{}, &SwarmOpts{
  1141  		OnInstanceRegistration: func() {
  1142  			client := ts.Client()
  1143  			req, err := http.NewRequest("GET", ts.URL, nil)
  1144  			req.Header.Set(rendezvous.HeaderID, id)
  1145  			req.Header.Set(rendezvous.HeaderToken, "test-token")
  1146  			req.Header.Set(rendezvous.HeaderHostname, "test-hostname")
  1147  			resp, err := client.Do(req)
  1148  			if err != nil {
  1149  				errCh <- fmt.Errorf("client.Do() = %s; want no error", err)
  1150  				return
  1151  			}
  1152  			if b, err := io.ReadAll(resp.Body); err != nil {
  1153  				errCh <- fmt.Errorf("io.ReadAll(body) = %b, %s, want no error", b, err)
  1154  				return
  1155  			}
  1156  			defer resp.Body.Close()
  1157  			if resp.StatusCode != 101 {
  1158  				errCh <- fmt.Errorf("resp.StatusCode  %d; want 101", resp.StatusCode)
  1159  				return
  1160  			}
  1161  		},
  1162  	}, false); err == nil || !strings.Contains(err.Error(), "revdial.Dialer closed") {
  1163  		errCh <- fmt.Errorf("startNewSwarmingTask() = bc, %s; want \"revdial.Dialer closed\" error", err)
  1164  	}
  1165  	select {
  1166  	case err := <-errCh:
  1167  		t.Fatal(err)
  1168  	default:
  1169  	}
  1170  }
  1171  
  1172  func mockSwarmClient() *swarmingtest.Client {
  1173  	return &swarmingtest.Client{
  1174  		NewTaskMock: func(context.Context, *swarmpb.NewTaskRequest) (*swarmpb.TaskRequestMetadataResponse, error) {
  1175  			panic("NewTask not implemented")
  1176  		},
  1177  		CountTasksMock: func(context.Context, float64, swarmpb.StateQuery, []string) (*swarmpb.TasksCount, error) {
  1178  			panic("CountTasks not implemented")
  1179  		},
  1180  		ListTasksMock: func(context.Context, int32, float64, swarmpb.StateQuery, []string) ([]*swarmpb.TaskResultResponse, error) {
  1181  			panic("ListTasks not implemented")
  1182  		},
  1183  		CancelTaskMock: func(context.Context, string, bool) (*swarmpb.CancelResponse, error) {
  1184  			panic("CancelTask not implemented")
  1185  		},
  1186  		TaskRequestMock: func(context.Context, string) (*swarmpb.TaskRequestResponse, error) {
  1187  			panic("TaskRequest not implemented")
  1188  		},
  1189  		TaskOutputMock: func(context.Context, string) (*swarmpb.TaskOutputResponse, error) {
  1190  			panic("TaskOutput not implemented")
  1191  		},
  1192  		TaskResultMock: func(context.Context, string, *swarming.TaskResultFields) (*swarmpb.TaskResultResponse, error) {
  1193  			panic("TaskResult not implemented")
  1194  		},
  1195  		CountBotsMock: func(context.Context, []*swarmpb.StringPair) (*swarmpb.BotsCount, error) {
  1196  			panic("CountBots not implemented")
  1197  		},
  1198  		ListBotsMock: func(context.Context, []*swarmpb.StringPair) ([]*swarmpb.BotInfo, error) {
  1199  			panic("ListBots not implemented")
  1200  		},
  1201  		DeleteBotMock: func(context.Context, string) (*swarmpb.DeleteResponse, error) { panic("TerminateBot not implemented") },
  1202  		TerminateBotMock: func(context.Context, string, string) (*swarmpb.TerminateResponse, error) {
  1203  			panic("TerminateBot not implemented")
  1204  		},
  1205  		ListBotTasksMock: func(context.Context, string, int32, float64, swarmpb.StateQuery) ([]*swarmpb.TaskResultResponse, error) {
  1206  			panic("ListBotTasks not implemented")
  1207  		},
  1208  		FilesFromCASMock: func(context.Context, string, *swarmpb.CASReference) ([]string, error) {
  1209  			panic("FilesFromCAS not implemented")
  1210  		},
  1211  	}
  1212  }
  1213  
  1214  func mockSwarmClientSimple() *swarmingtest.Client {
  1215  	msc := mockSwarmClient()
  1216  	msc.NewTaskMock = func(_ context.Context, req *swarmpb.NewTaskRequest) (*swarmpb.TaskRequestMetadataResponse, error) {
  1217  		taskID := uuid.New().String()
  1218  		return &swarmpb.TaskRequestMetadataResponse{
  1219  			TaskId: taskID,
  1220  			Request: &swarmpb.TaskRequestResponse{
  1221  				TaskId: taskID,
  1222  				Name:   req.Name,
  1223  			},
  1224  		}, nil
  1225  	}
  1226  	msc.TaskResultMock = func(_ context.Context, taskID string, _ *swarming.TaskResultFields) (*swarmpb.TaskResultResponse, error) {
  1227  		return &swarmpb.TaskResultResponse{
  1228  			TaskId: taskID,
  1229  			State:  swarmpb.TaskState_RUNNING,
  1230  		}, nil
  1231  	}
  1232  	return msc
  1233  }
  1234  
  1235  func mustCreateSwarmingInstance(t *testing.T, client protos.GomoteServiceClient, iap access.IAPFields) string {
  1236  	req := &protos.CreateInstanceRequest{
  1237  		BuilderType: "gotip-linux-amd64-boringcrypto",
  1238  	}
  1239  	stream, err := client.CreateInstance(access.FakeContextWithOutgoingIAPAuth(context.Background(), iap), req)
  1240  	if err != nil {
  1241  		t.Fatalf("client.CreateInstance(ctx, %v) = %v,  %s; want no error", req, stream, err)
  1242  	}
  1243  	var updateComplete bool
  1244  	var gomoteID string
  1245  	for {
  1246  		update, err := stream.Recv()
  1247  		if err == io.EOF && !updateComplete {
  1248  			t.Fatal("stream.Recv = stream, io.EOF; want no EOF")
  1249  		}
  1250  		if err == io.EOF {
  1251  			break
  1252  		}
  1253  		if err != nil {
  1254  			t.Fatalf("stream.Recv() = nil, %s; want no error", err)
  1255  		}
  1256  		if update.GetStatus() == protos.CreateInstanceResponse_COMPLETE {
  1257  			gomoteID = update.Instance.GetGomoteId()
  1258  			updateComplete = true
  1259  		}
  1260  	}
  1261  	return gomoteID
  1262  }
  1263  
  1264  type FakeBuildersClient struct{}
  1265  
  1266  func (fbc *FakeBuildersClient) GetBuilder(ctx context.Context, in *buildbucketpb.GetBuilderRequest, opts ...grpc.CallOption) (*buildbucketpb.BuilderItem, error) {
  1267  	builders := map[string]bool{
  1268  		"gotip-linux-amd64-boringcrypto-test_only": true,
  1269  	}
  1270  	name := in.GetId().GetBuilder()
  1271  	_, ok := builders[name]
  1272  	if !ok {
  1273  		return nil, errors.New("builder type not found")
  1274  	}
  1275  	return &buildbucketpb.BuilderItem{
  1276  		Id: &buildbucketpb.BuilderID{
  1277  			Project: "golang",
  1278  			Bucket:  "ci-workers",
  1279  			Builder: name,
  1280  		},
  1281  		Config: &buildbucketpb.BuilderConfig{
  1282  			Name: name,
  1283  			Dimensions: []string{
  1284  				"cipd_platform:linux-amd64",
  1285  			},
  1286  		},
  1287  	}, nil
  1288  }
  1289  
  1290  func (fbc *FakeBuildersClient) ListBuilders(ctx context.Context, in *buildbucketpb.ListBuildersRequest, opts ...grpc.CallOption) (*buildbucketpb.ListBuildersResponse, error) {
  1291  	makeBuilderItem := func(bucket string, builders ...string) []*buildbucketpb.BuilderItem {
  1292  		out := make([]*buildbucketpb.BuilderItem, 0, len(builders))
  1293  		for _, b := range builders {
  1294  			out = append(out, &buildbucketpb.BuilderItem{
  1295  				Id: &buildbucketpb.BuilderID{
  1296  					Project: "golang",
  1297  					Bucket:  bucket,
  1298  					Builder: b,
  1299  				},
  1300  				Config: &buildbucketpb.BuilderConfig{
  1301  					Name: b,
  1302  					Dimensions: []string{
  1303  						"cipd_platform:linux-amd64",
  1304  					},
  1305  					Properties: `{"mode": 0, "bootstrap_version":"latest"}`,
  1306  				},
  1307  			})
  1308  		}
  1309  		return out
  1310  	}
  1311  	var builders []*buildbucketpb.BuilderItem
  1312  	switch bucket := in.GetBucket(); bucket {
  1313  	case "ci-workers":
  1314  		builders = makeBuilderItem(bucket, "gotip-linux-amd64-boringcrypto", "gotip-linux-amd64-boringcrypto-test_only")
  1315  	case "ci":
  1316  		builders = makeBuilderItem(bucket, "gotip-linux-arm", "gotip-linux-amd64")
  1317  	default:
  1318  		builders = []*buildbucketpb.BuilderItem{}
  1319  	}
  1320  	out := &buildbucketpb.ListBuildersResponse{
  1321  		Builders:      builders,
  1322  		NextPageToken: "",
  1323  	}
  1324  	return out, nil
  1325  }