go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/logdog/appengine/coordinator/endpoints/services/taskQueue_test.go (about)

     1  // Copyright 2019 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 services
    16  
    17  import (
    18  	"math/rand"
    19  	"sort"
    20  	"testing"
    21  	"time"
    22  
    23  	"google.golang.org/protobuf/types/known/durationpb"
    24  
    25  	"go.chromium.org/luci/common/data/rand/mathrand"
    26  	"go.chromium.org/luci/gae/service/taskqueue"
    27  	logdog "go.chromium.org/luci/logdog/api/endpoints/coordinator/services/v1"
    28  	ct "go.chromium.org/luci/logdog/appengine/coordinator/coordinatorTest"
    29  
    30  	. "github.com/smartystreets/goconvey/convey"
    31  )
    32  
    33  func TestTaskQueue(t *testing.T) {
    34  	Convey(`With a testing configuration`, t, func() {
    35  		c, env := ct.Install()
    36  		c = mathrand.Set(c, rand.New(rand.NewSource(1234)))
    37  
    38  		// By default, the testing user is a service.
    39  		env.ActAsService()
    40  		svr := New(ServerSettings{NumQueues: 2})
    41  
    42  		// The testable TQ object.
    43  		ts := taskqueue.GetTestable(c)
    44  		ts.CreatePullQueue(RawArchiveQueueName(0))
    45  		ts.CreatePullQueue(RawArchiveQueueName(1))
    46  
    47  		mustAddTask := func(t *logdog.ArchiveTask) {
    48  			task, err := tqTask(t)
    49  			So(err, ShouldBeNil)
    50  
    51  			queueName, _ := svr.(*logdog.DecoratedServices).Service.(*server).getNextArchiveQueueName(c)
    52  
    53  			So(taskqueue.Add(c, queueName, task), ShouldBeNil)
    54  		}
    55  
    56  		Convey(`Lease a task with empty taskqueue`, func() {
    57  			tasks, err := svr.LeaseArchiveTasks(c, &logdog.LeaseRequest{
    58  				MaxTasks:  10,
    59  				LeaseTime: durationpb.New(10 * time.Minute),
    60  			})
    61  			So(err, ShouldBeNil)
    62  			So(len(tasks.Tasks), ShouldEqual, 0)
    63  		})
    64  		task1 := &logdog.ArchiveTask{Project: "foo", Id: "deadbeef1", Realm: "foo:bar"}
    65  		task2 := &logdog.ArchiveTask{Project: "foo", Id: "deadbeef2", Realm: "foo:bar"}
    66  		task3 := &logdog.ArchiveTask{Project: "foo", Id: "deadbeef3", Realm: "foo:bar"}
    67  		task4 := &logdog.ArchiveTask{Project: "foo", Id: "deadbeef4", Realm: "foo:bar"}
    68  
    69  		Convey(`Two tasks`, func() {
    70  			mustAddTask(task1)
    71  			mustAddTask(task2)
    72  
    73  			var leasedTasks []*logdog.ArchiveTask
    74  			// We have to retry a couple times to collect all the tasks b/c
    75  			// randomness. This is sensitive to the math.NewSource() value above.
    76  			for i := 0; i < 3; i++ {
    77  				tasks, err := svr.LeaseArchiveTasks(c, &logdog.LeaseRequest{
    78  					MaxTasks:  10,
    79  					LeaseTime: durationpb.New(10 * time.Minute),
    80  				})
    81  				So(err, ShouldBeNil)
    82  				leasedTasks = append(leasedTasks, tasks.Tasks...)
    83  			}
    84  
    85  			So(len(leasedTasks), ShouldEqual, 2)
    86  			// No order is guaranteed in the returned slice
    87  			sort.Slice(leasedTasks, func(i, j int) bool {
    88  				return leasedTasks[i].Id < leasedTasks[j].Id
    89  			})
    90  			So(leasedTasks[0].Project, ShouldEqual, "foo")
    91  			So(leasedTasks[0].Id, ShouldEqual, task1.Id)
    92  			So(leasedTasks[1].Id, ShouldEqual, task2.Id)
    93  			Convey(`And delete one of the tasks`, func() {
    94  				_, err := svr.DeleteArchiveTasks(c, &logdog.DeleteRequest{
    95  					// leasedTasks[0] has TaskName filled, but task1 does not.
    96  					Tasks: []*logdog.ArchiveTask{leasedTasks[0]},
    97  				})
    98  				So(err, ShouldBeNil)
    99  
   100  				// there should only be one scheduled task remaining.
   101  				numScheduled := 0
   102  				for _, tasks := range ts.GetScheduledTasks() {
   103  					numScheduled += len(tasks)
   104  				}
   105  
   106  				So(numScheduled, ShouldEqual, 1)
   107  			})
   108  		})
   109  
   110  		Convey(`Many tasks`, func() {
   111  			mustAddTask(task1)
   112  			mustAddTask(task2)
   113  			mustAddTask(task3)
   114  			mustAddTask(task4)
   115  
   116  			var leasedTasks []*logdog.ArchiveTask
   117  			// We have to retry a couple times to collect all the tasks b/c
   118  			// randomness. This is sensitive to the math.NewSource() value above.
   119  			for i := 0; i < 3; i++ {
   120  				tasks, err := svr.LeaseArchiveTasks(c, &logdog.LeaseRequest{
   121  					MaxTasks:  1,
   122  					LeaseTime: durationpb.New(10 * time.Minute),
   123  				})
   124  				So(err, ShouldBeNil)
   125  				leasedTasks = append(leasedTasks, tasks.Tasks...)
   126  			}
   127  
   128  			So(len(leasedTasks), ShouldEqual, 3)
   129  		})
   130  
   131  	})
   132  }