go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/tq/internal/submit_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 internal
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"sort"
    21  	"sync"
    22  	"testing"
    23  
    24  	taskspb "cloud.google.com/go/cloudtasks/apiv2/cloudtaskspb"
    25  	"google.golang.org/grpc/codes"
    26  	"google.golang.org/grpc/status"
    27  
    28  	"go.chromium.org/luci/common/data/stringset"
    29  
    30  	"go.chromium.org/luci/server/tq/internal/reminder"
    31  	"go.chromium.org/luci/server/tq/internal/testutil"
    32  
    33  	. "github.com/smartystreets/goconvey/convey"
    34  )
    35  
    36  func TestSubmit(t *testing.T) {
    37  	t.Parallel()
    38  
    39  	Convey("With mocks", t, func() {
    40  		ctx := context.Background()
    41  		db := testutil.FakeDB{}
    42  		sub := submitter{}
    43  		ctx = db.Inject(ctx)
    44  
    45  		makeRem := func(id string) *reminder.Reminder {
    46  			r := &reminder.Reminder{ID: id}
    47  			r.AttachPayload(&reminder.Payload{
    48  				CreateTaskRequest: &taskspb.CreateTaskRequest{
    49  					Parent: id + " body",
    50  				},
    51  			})
    52  			So(db.SaveReminder(ctx, r), ShouldBeNil)
    53  			return r
    54  		}
    55  
    56  		call := func(r *reminder.Reminder) error {
    57  			return SubmitFromReminder(ctx, &sub, &db, r, TxnPathHappy)
    58  		}
    59  
    60  		r := makeRem("rem")
    61  		So(db.AllReminders(), ShouldHaveLength, 1)
    62  
    63  		Convey("Success, have payload deserialized", func() {
    64  			So(call(r), ShouldBeNil)
    65  			So(db.AllReminders(), ShouldHaveLength, 0)
    66  			So(sub.req, ShouldHaveLength, 1)
    67  			So(sub.req[0].CreateTaskRequest.Parent, ShouldEqual, "rem body")
    68  		})
    69  
    70  		Convey("Success, from raw reminder payload", func() {
    71  			So(call(r.DropPayload()), ShouldBeNil)
    72  			So(db.AllReminders(), ShouldHaveLength, 0)
    73  			So(sub.req, ShouldHaveLength, 1)
    74  			So(sub.req[0].CreateTaskRequest.Parent, ShouldEqual, "rem body")
    75  		})
    76  
    77  		Convey("Transient err", func() {
    78  			sub.err = status.Errorf(codes.Internal, "boo")
    79  			So(call(r), ShouldNotBeNil)
    80  			So(db.AllReminders(), ShouldHaveLength, 1) // kept it
    81  		})
    82  
    83  		Convey("Fatal err", func() {
    84  			sub.err = status.Errorf(codes.PermissionDenied, "boo")
    85  			So(call(r), ShouldNotBeNil)
    86  			So(db.AllReminders(), ShouldHaveLength, 0) // deleted it
    87  		})
    88  
    89  		Convey("Batch", func() {
    90  			for i := 0; i < 5; i++ {
    91  				makeRem(fmt.Sprintf("more-%d", i))
    92  			}
    93  
    94  			batch := db.AllReminders()
    95  			for _, r := range batch {
    96  				r.RawPayload = nil
    97  			}
    98  			So(batch, ShouldHaveLength, 6)
    99  
   100  			// Reject `rem`, keep only `more-...`.
   101  			sub.ban = stringset.NewFromSlice("rem body")
   102  
   103  			n, err := SubmitBatch(ctx, &sub, &db, batch)
   104  			So(err, ShouldNotBeNil) // had failing requests
   105  			So(n, ShouldEqual, 5)   // still submitted 5 tasks
   106  
   107  			// Verify they had correct payloads.
   108  			var bodies []string
   109  			for _, r := range sub.req {
   110  				bodies = append(bodies, r.CreateTaskRequest.Parent)
   111  			}
   112  			sort.Strings(bodies)
   113  			So(bodies, ShouldResemble, []string{
   114  				"more-0 body", "more-1 body", "more-2 body", "more-3 body", "more-4 body",
   115  			})
   116  
   117  			// All reminders are deleted, even the one that matches the failed task.
   118  			So(db.AllReminders(), ShouldHaveLength, 0)
   119  		})
   120  	})
   121  }
   122  
   123  type submitter struct {
   124  	m   sync.Mutex
   125  	ban stringset.Set
   126  	err error
   127  	req []*reminder.Payload
   128  }
   129  
   130  func (s *submitter) Submit(ctx context.Context, req *reminder.Payload) error {
   131  	s.m.Lock()
   132  	defer s.m.Unlock()
   133  	if s.ban.Has(req.CreateTaskRequest.Parent) {
   134  		return status.Errorf(codes.PermissionDenied, "boom")
   135  	}
   136  	s.req = append(s.req, req)
   137  	return s.err
   138  }