tasks/common/quotas/priority_rate_limiter_impl_test.go at release/v1.30.x

mirror of https://github.com/hanzoai/tasks.git synced 2026-10-04 02:04:03 +00:00

Files

T

Alex ShtinandGitHub 8e8f1b0f7c Remove license header from every file (#7689 )

## What changed?
<!-- Describe what has changed in this PR -->
Remove license header from every file. Because it is really hard to
follow in this PR here is the summary:
1. License header is removed from all `.go` and `.proto` files
:fireworks::fireworks:🎆.
2. `LICENSE` file in the root directory has only Temporal and Uber
copyrights.
3. 5 other `LICENSE` files added to the packages which have copyrights
different from Temporal and Uber: Datadog, Xargin, "Mat Ryer, Tyler
Bunnell and contributors".
4. `license_file` flag is removed from all code generation tools.
5. `copyright_file` flag is removed from `go:generate mockgen`
directive.
6. All copyright related targets are removed from `Makefile`.
7. Updated Temporal copyright year to 2025 everywhere.

## Why?
<!-- Tell your future self why have you made these changes -->
I double checked with legal department that it is not needed to have
license header in every file. One file per repo is enough. I put all
copyrights to the root `LICENSE` file and removed header from all other
files. Also updated tools and `Makefile`.

2025-05-01 18:50:21 -07:00

411 lines

12 KiB

Go

Raw Permalink Blame History

 package quotas
 import (
 	"context"
 	"testing"
 	"time"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 	"go.uber.org/mock/gomock"
 )
 type (
 	priorityStageRateLimiterSuite struct {
 		suite.Suite
 		*require.Assertions
 		controller              *gomock.Controller
 		highPriorityRateLimiter *MockRateLimiter
 		lowPriorityRateLimiter  *MockRateLimiter
 		highPriorityReservation *MockReservation
 		lowPriorityReservation  *MockReservation
 		highPriorityAPIName     string
 		lowPriorityAPIName      string
 		rateLimiter *PriorityRateLimiterImpl
 	}
 )
 func TestPriorityStageRateLimiterSuite(t *testing.T) {
 	s := new(priorityStageRateLimiterSuite)
 	suite.Run(t, s)
 }
 func (s *priorityStageRateLimiterSuite) SetupSuite() {
 }
 func (s *priorityStageRateLimiterSuite) TearDownSuite() {
 }
 func (s *priorityStageRateLimiterSuite) SetupTest() {
 	s.Assertions = require.New(s.T())
 	s.controller = gomock.NewController(s.T())
 	s.highPriorityRateLimiter = NewMockRateLimiter(s.controller)
 	s.lowPriorityRateLimiter = NewMockRateLimiter(s.controller)
 	s.highPriorityReservation = NewMockReservation(s.controller)
 	s.lowPriorityReservation = NewMockReservation(s.controller)
 	s.highPriorityAPIName = "high-priority"
 	s.lowPriorityAPIName = "low-priority"
 	apiToPriority := map[string]int{
 		s.highPriorityAPIName: 0,
 		s.lowPriorityAPIName:  2,
 	}
 	priorityToRateLimiters := map[int]RequestRateLimiter{
 : NewRequestRateLimiterAdapter(s.highPriorityRateLimiter),
 : NewRequestRateLimiterAdapter(s.lowPriorityRateLimiter),
 	}
 	s.rateLimiter = NewPriorityRateLimiter(func(req Request) int {
 		return apiToPriority[req.API]
 	}, priorityToRateLimiters)
 }
 func (s *priorityStageRateLimiterSuite) TearDownTest() {
 	s.controller.Finish()
 }
 func (s *priorityStageRateLimiterSuite) TestAllow_HighPriority_Allow() {
 	now := time.Now()
 	token := 1
 	req := Request{
 		API:    s.highPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	s.highPriorityRateLimiter.EXPECT().AllowN(now, token).Return(true)
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(now, token).Return(s.lowPriorityReservation)
 	allow := s.rateLimiter.Allow(now, req)
 	s.True(allow)
 }
 func (s *priorityStageRateLimiterSuite) TestAllow_HighPriority_Disallow() {
 	now := time.Now()
 	token := 1
 	req := Request{
 		API:    s.highPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	s.highPriorityRateLimiter.EXPECT().AllowN(now, token).Return(false)
 	allow := s.rateLimiter.Allow(now, req)
 	s.False(allow)
 }
 func (s *priorityStageRateLimiterSuite) TestAllow_LowPriority_Allow() {
 	now := time.Now()
 	token := 1
 	req := Request{
 		API:    s.lowPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	s.lowPriorityRateLimiter.EXPECT().AllowN(now, token).Return(true)
 	allow := s.rateLimiter.Allow(now, req)
 	s.True(allow)
 }
 func (s *priorityStageRateLimiterSuite) TestAllow_LowPriority_Disallow() {
 	now := time.Now()
 	token := 1
 	req := Request{
 		API:    s.lowPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	s.lowPriorityRateLimiter.EXPECT().AllowN(now, token).Return(false)
 	allow := s.rateLimiter.Allow(now, req)
 	s.False(allow)
 }
 func (s *priorityStageRateLimiterSuite) TestReserve_HighPriority_OK() {
 	now := time.Now()
 	token := 1
 	req := Request{
 		API:    s.highPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	s.highPriorityReservation.EXPECT().OK().Return(true)
 	s.highPriorityRateLimiter.EXPECT().ReserveN(now, token).Return(s.highPriorityReservation)
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(now, token).Return(s.lowPriorityReservation)
 	reservation := s.rateLimiter.Reserve(now, req)
 	s.Equal(NewPriorityReservation(
 		s.highPriorityReservation,
 		[]Reservation{s.lowPriorityReservation},
 	), reservation)
 }
 func (s *priorityStageRateLimiterSuite) TestReserve_HighPriority_NotOK() {
 	now := time.Now()
 	token := 1
 	req := Request{
 		API:    s.highPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	s.highPriorityReservation.EXPECT().OK().Return(false)
 	s.highPriorityRateLimiter.EXPECT().ReserveN(now, token).Return(s.highPriorityReservation)
 	reservation := s.rateLimiter.Reserve(now, req)
 	s.Equal(s.highPriorityReservation, reservation)
 }
 func (s *priorityStageRateLimiterSuite) TestReserve_LowPriority_OK() {
 	now := time.Now()
 	token := 1
 	req := Request{
 		API:    s.lowPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	s.lowPriorityReservation.EXPECT().OK().Return(true)
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(now, token).Return(s.lowPriorityReservation)
 	reservation := s.rateLimiter.Reserve(now, req)
 	s.Equal(NewPriorityReservation(
 		s.lowPriorityReservation,
 		[]Reservation{},
 	), reservation)
 }
 func (s *priorityStageRateLimiterSuite) TestReserve_LowPriority_NotOK() {
 	now := time.Now()
 	token := 1
 	req := Request{
 		API:    s.lowPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	s.lowPriorityReservation.EXPECT().OK().Return(false)
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(now, token).Return(s.lowPriorityReservation)
 	reservation := s.rateLimiter.Reserve(now, req)
 	s.Equal(s.lowPriorityReservation, reservation)
 }
 func (s *priorityStageRateLimiterSuite) TestWait_HighPriority_AlreadyExpired() {
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
 	token := 1
 	req := Request{
 		API:    s.highPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	err := s.rateLimiter.Wait(ctx, req)
 	s.Error(err)
 }
 func (s *priorityStageRateLimiterSuite) TestWait_LowPriority_AlreadyExpired() {
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
 	token := 1
 	req := Request{
 		API:    s.lowPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	err := s.rateLimiter.Wait(ctx, req)
 	s.Error(err)
 }
 func (s *priorityStageRateLimiterSuite) TestWait_HighPriority_NotExpired_WithExpiration_Error() {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
 	token := 1
 	req := Request{
 		API:    s.highPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	highPriorityReservationDelay := 2 * time.Second
 	s.highPriorityReservation.EXPECT().DelayFrom(gomock.Any()).Return(highPriorityReservationDelay)
 	s.highPriorityReservation.EXPECT().CancelAt(gomock.Any())
 	s.lowPriorityReservation.EXPECT().CancelAt(gomock.Any())
 	s.highPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.highPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.highPriorityReservation)
 	s.lowPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.lowPriorityReservation)
 	err := s.rateLimiter.Wait(ctx, req)
 	s.Error(err)
 }
 func (s *priorityStageRateLimiterSuite) TestWait_LowPriority_NotExpired_WithExpiration_Error() {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
 	token := 1
 	req := Request{
 		API:    s.lowPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	lowPriorityReservationDelay := 3 * time.Second
 	s.lowPriorityReservation.EXPECT().DelayFrom(gomock.Any()).Return(lowPriorityReservationDelay)
 	s.lowPriorityReservation.EXPECT().CancelAt(gomock.Any())
 	s.lowPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.lowPriorityReservation)
 	err := s.rateLimiter.Wait(ctx, req)
 	s.Error(err)
 }
 func (s *priorityStageRateLimiterSuite) TestWait_HighPriority_NotExpired_WithExpiration_Cancelled() {
 	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Second)
 	token := 1
 	req := Request{
 		API:    s.highPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	go func() {
 		time.Sleep(4 * time.Second)
 		cancel()
 	}()
 	highPriorityReservationDelay := 20 * time.Second
 	s.highPriorityReservation.EXPECT().DelayFrom(gomock.Any()).Return(highPriorityReservationDelay)
 	s.highPriorityReservation.EXPECT().CancelAt(gomock.Any())
 	s.lowPriorityReservation.EXPECT().CancelAt(gomock.Any())
 	s.highPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.highPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.highPriorityReservation)
 	s.lowPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.lowPriorityReservation)
 	err := s.rateLimiter.Wait(ctx, req)
 	s.Error(err)
 }
 func (s *priorityStageRateLimiterSuite) TestWait_LowPriority_NotExpired_WithExpiration_Cancelled() {
 	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Second)
 	token := 1
 	req := Request{
 		API:    s.lowPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	go func() {
 		time.Sleep(4 * time.Second)
 		cancel()
 	}()
 	lowPriorityReservationDelay := 30 * time.Second
 	s.lowPriorityReservation.EXPECT().DelayFrom(gomock.Any()).Return(lowPriorityReservationDelay)
 	s.lowPriorityReservation.EXPECT().CancelAt(gomock.Any())
 	s.lowPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.lowPriorityReservation)
 	err := s.rateLimiter.Wait(ctx, req)
 	s.Error(err)
 }
 func (s *priorityStageRateLimiterSuite) TestWait_HighPriority_NotExpired_WithExpiration_NoError() {
 	ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second)
 	defer cancel()
 	token := 1
 	req := Request{
 		API:    s.highPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	highPriorityReservationDelay := 2 * time.Second
 	s.highPriorityReservation.EXPECT().DelayFrom(gomock.Any()).Return(highPriorityReservationDelay)
 	s.highPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.highPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.highPriorityReservation)
 	s.lowPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.lowPriorityReservation)
 	err := s.rateLimiter.Wait(ctx, req)
 	s.NoError(err)
 }
 func (s *priorityStageRateLimiterSuite) TestWait_LowPriority_NotExpired_WithExpiration_NoError() {
 	ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second)
 	defer cancel()
 	token := 1
 	req := Request{
 		API:    s.lowPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	lowPriorityReservationDelay := 2 * time.Second
 	s.lowPriorityReservation.EXPECT().DelayFrom(gomock.Any()).Return(lowPriorityReservationDelay)
 	s.lowPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.lowPriorityReservation)
 	err := s.rateLimiter.Wait(ctx, req)
 	s.NoError(err)
 }
 func (s *priorityStageRateLimiterSuite) TestWait_HighPriority_NotExpired_WithoutExpiration() {
 	ctx := context.Background()
 	token := 1
 	req := Request{
 		API:    s.highPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	highPriorityReservationDelay := 2 * time.Second
 	s.highPriorityReservation.EXPECT().DelayFrom(gomock.Any()).Return(highPriorityReservationDelay)
 	s.highPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.highPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.highPriorityReservation)
 	s.lowPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.lowPriorityReservation)
 	err := s.rateLimiter.Wait(ctx, req)
 	s.NoError(err)
 }
 func (s *priorityStageRateLimiterSuite) TestWait_LowPriority_NotExpired_WithoutExpiration() {
 	ctx := context.Background()
 	token := 1
 	req := Request{
 		API:    s.lowPriorityAPIName,
 		Token:  token,
 		Caller: "",
 	}
 	lowPriorityReservationDelay := 3 * time.Second
 	s.lowPriorityReservation.EXPECT().DelayFrom(gomock.Any()).Return(lowPriorityReservationDelay)
 	s.lowPriorityReservation.EXPECT().OK().Return(true).AnyTimes()
 	s.lowPriorityRateLimiter.EXPECT().ReserveN(gomock.Any(), token).Return(s.lowPriorityReservation)
 	err := s.rateLimiter.Wait(ctx, req)
 	s.NoError(err)
 }

411 lines 12 KiB Go Raw Permalink Blame History

411 lines

12 KiB

Go

Raw Permalink Blame History