go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/resultdb/internal/integrationtests/testclient_test.go (about)

     1  // Copyright 2020 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package integrationtests
    16  
    17  import (
    18  	"context"
    19  	"time"
    20  
    21  	"google.golang.org/genproto/protobuf/field_mask"
    22  	"google.golang.org/grpc"
    23  	"google.golang.org/grpc/metadata"
    24  
    25  	"go.chromium.org/luci/resultdb/internal/services/recorder"
    26  	"go.chromium.org/luci/resultdb/pbutil"
    27  	pb "go.chromium.org/luci/resultdb/proto/v1"
    28  
    29  	. "github.com/smartystreets/goconvey/convey"
    30  )
    31  
    32  // testClient is a convenient resultdb client, to keep tests simple.
    33  // Asserts that all requests succeed.
    34  // Memorizes update tokens of invocations it created.
    35  type testClient struct {
    36  	app *testApp
    37  
    38  	updateTokens map[string]string
    39  }
    40  
    41  func (c *testClient) CreateInvocation(ctx context.Context, id string) {
    42  	md := metadata.MD{}
    43  	req := &pb.CreateInvocationRequest{InvocationId: id, Invocation: &pb.Invocation{Realm: "testproject:testrealm"}}
    44  	inv, err := c.app.Recorder.CreateInvocation(ctx, req, grpc.Header(&md))
    45  	So(err, ShouldBeNil)
    46  	So(md.Get(pb.UpdateTokenMetadataKey), ShouldHaveLength, 1)
    47  
    48  	if c.updateTokens == nil {
    49  		c.updateTokens = map[string]string{}
    50  	}
    51  	c.updateTokens[inv.Name] = md.Get(pb.UpdateTokenMetadataKey)[0]
    52  }
    53  
    54  func (c *testClient) withUpdateTokenFor(ctx context.Context, invocation string) context.Context {
    55  	return metadata.AppendToOutgoingContext(ctx, pb.UpdateTokenMetadataKey, c.updateTokens[invocation])
    56  }
    57  
    58  func (c *testClient) GetState(ctx context.Context, name string) pb.Invocation_State {
    59  	inv, err := c.app.ResultDB.GetInvocation(ctx, &pb.GetInvocationRequest{Name: name})
    60  	So(err, ShouldBeNil)
    61  	return inv.State
    62  }
    63  
    64  func (c *testClient) Include(ctx context.Context, including, included string) {
    65  	ctx = c.withUpdateTokenFor(ctx, including)
    66  	_, err := c.app.Recorder.UpdateIncludedInvocations(ctx, &pb.UpdateIncludedInvocationsRequest{
    67  		IncludingInvocation: including,
    68  		AddInvocations:      []string{included},
    69  	})
    70  	So(err, ShouldBeNil)
    71  }
    72  
    73  func (c *testClient) FinalizeInvocation(ctx context.Context, name string) {
    74  	ctx = c.withUpdateTokenFor(ctx, name)
    75  	_, err := c.app.Recorder.FinalizeInvocation(ctx, &pb.FinalizeInvocationRequest{Name: name})
    76  	So(err, ShouldBeNil)
    77  }
    78  
    79  // MakeInvocationOverdue uses a magic constant to set an invocation's deadline
    80  // sometime in the past.
    81  func (c *testClient) MakeInvocationOverdue(ctx context.Context, name string) {
    82  	ctx = c.withUpdateTokenFor(ctx, name)
    83  	_, err := c.app.Recorder.UpdateInvocation(ctx, &pb.UpdateInvocationRequest{
    84  		Invocation: &pb.Invocation{
    85  			Name:     name,
    86  			Deadline: pbutil.MustTimestampProto(time.Unix(recorder.TestMagicOverdueDeadlineUnixSecs, 0).UTC()),
    87  		},
    88  		UpdateMask: &field_mask.FieldMask{
    89  			Paths: []string{"deadline"},
    90  		},
    91  	})
    92  	So(err, ShouldBeNil)
    93  }