go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/tq/distributed_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 tq
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"testing"
    21  
    22  	"go.chromium.org/luci/server/tq/internal/tqpb"
    23  	"go.chromium.org/luci/server/tq/tqtesting"
    24  
    25  	. "github.com/smartystreets/goconvey/convey"
    26  	. "go.chromium.org/luci/common/testing/assertions"
    27  )
    28  
    29  func TestSweepRouting(t *testing.T) {
    30  	t.Parallel()
    31  
    32  	Convey("Works", t, func() {
    33  		ctx := context.Background()
    34  
    35  		disp := Dispatcher{}
    36  		ctx, sched := TestingContext(ctx, &disp)
    37  
    38  		mu := sync.Mutex{}
    39  		calls := []*tqpb.SweepTask{}
    40  
    41  		enqueue := sweepTaskRouting(&disp,
    42  			DistributedSweeperOptions{TaskQueue: "zzz"},
    43  			func(_ context.Context, task *tqpb.SweepTask) error {
    44  				mu.Lock()
    45  				calls = append(calls, task)
    46  				mu.Unlock()
    47  				return nil
    48  			},
    49  		)
    50  
    51  		submitted := &tqpb.SweepTask{ShardCount: 123}
    52  		enqueue(ctx, submitted)
    53  
    54  		sched.Run(ctx, tqtesting.StopWhenDrained())
    55  		So(calls, ShouldHaveLength, 1)
    56  		So(calls[0], ShouldResembleProto, submitted)
    57  	})
    58  }