go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/cv/internal/rpc/admin/dsmapper_test.go (about)

     1  // Copyright 2021 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 admin
    16  
    17  import (
    18  	"context"
    19  	"testing"
    20  
    21  	"go.chromium.org/luci/gae/service/datastore"
    22  	"go.chromium.org/luci/server/auth"
    23  	"go.chromium.org/luci/server/auth/authtest"
    24  	"go.chromium.org/luci/server/dsmapper"
    25  	"go.chromium.org/luci/server/dsmapper/dsmapperpb"
    26  	"go.chromium.org/luci/server/tq/tqtesting"
    27  
    28  	"go.chromium.org/luci/cv/internal/cvtesting"
    29  	adminpb "go.chromium.org/luci/cv/internal/rpc/admin/api"
    30  
    31  	. "github.com/smartystreets/goconvey/convey"
    32  	. "go.chromium.org/luci/common/testing/assertions"
    33  )
    34  
    35  func TestDSMapperServer(t *testing.T) {
    36  	t.Parallel()
    37  
    38  	Convey("dsmapper job lifecycle", t, func() {
    39  		ct := cvtesting.Test{}
    40  		ctx, cancel := ct.SetUp(t)
    41  		defer cancel()
    42  
    43  		d := dsMapper{
    44  			ctrl: &dsmapper.Controller{
    45  				ControlQueue: "default",
    46  				MapperQueue:  "default",
    47  			},
    48  		}
    49  		d.ctrl.Install(ct.TQDispatcher)
    50  		d.register(
    51  			&dsmapper.JobConfig{
    52  				Mapper: "upgrade-something",
    53  				Query: dsmapper.Query{
    54  					Kind: "SomethingUniqueTest",
    55  				},
    56  				ShardCount: 2,
    57  				PageSize:   16,
    58  			},
    59  			func(context.Context, *dsmapper.Job, int) (dsmapper.Mapper, error) {
    60  				return func(ctx context.Context, keys []*datastore.Key) error {
    61  					// In prod, this updates the entities with the given keys.
    62  					return nil
    63  				}, nil
    64  			},
    65  		)
    66  		a := AdminServer{dsmapper: &d}
    67  
    68  		Convey("without access", func() {
    69  			ctx = auth.WithState(ctx, &authtest.FakeState{
    70  				Identity: "anonymous:anonymous",
    71  			})
    72  			_, err := a.DSMLaunchJob(ctx, &adminpb.DSMLaunchJobRequest{Name: "upgrade-something"})
    73  			So(err, ShouldBeRPCPermissionDenied)
    74  			_, err = a.DSMGetJob(ctx, &adminpb.DSMJobID{Id: 1})
    75  			So(err, ShouldBeRPCPermissionDenied)
    76  			_, err = a.DSMAbortJob(ctx, &adminpb.DSMJobID{Id: 1})
    77  			So(err, ShouldBeRPCPermissionDenied)
    78  		})
    79  
    80  		Convey("with access", func() {
    81  			ctx = auth.WithState(ctx, &authtest.FakeState{
    82  				Identity:       "user:admin@example.com",
    83  				IdentityGroups: []string{allowGroup},
    84  			})
    85  
    86  			jobID, err := a.DSMLaunchJob(ctx, &adminpb.DSMLaunchJobRequest{Name: "upgrade-something"})
    87  			So(err, ShouldBeNil)
    88  
    89  			job, err := a.DSMGetJob(ctx, jobID)
    90  			So(err, ShouldBeNil)
    91  			So(job.GetName(), ShouldResemble, "upgrade-something")
    92  			So(job.GetInfo().GetState(), ShouldEqual, dsmapperpb.State_STARTING)
    93  
    94  			Convey("SUCCESS", func() {
    95  				ct.TQ.Run(ctx, tqtesting.StopWhenDrained())
    96  
    97  				job, err = a.DSMGetJob(ctx, jobID)
    98  				So(err, ShouldBeNil)
    99  				So(job.GetName(), ShouldResemble, "upgrade-something")
   100  				So(job.GetInfo().GetState(), ShouldEqual, dsmapperpb.State_SUCCESS)
   101  			})
   102  			Convey("Abort", func() {
   103  				_, err = a.DSMAbortJob(ctx, jobID)
   104  				So(err, ShouldBeNil)
   105  				ct.TQ.Run(ctx, tqtesting.StopWhenDrained())
   106  				// This fails with:
   107  				//   "broken state, no ShardList entity for job 1"
   108  				// which is probably because job was aborted right after launching.
   109  				// job, err = d.DSMGetJob(ctx, jobID)
   110  				// So(err, ShouldBeNil)
   111  			})
   112  		})
   113  	})
   114  }