You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
64 lines
1.2 KiB
64 lines
1.2 KiB
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
|
|
"github.com/go-kratos/aegis/ratelimit"
|
|
)
|
|
|
|
type (
|
|
ratelimitMock struct {
|
|
reached bool
|
|
}
|
|
ratelimitReachedMock struct {
|
|
reached bool
|
|
}
|
|
)
|
|
|
|
func (r *ratelimitMock) Allow() (ratelimit.DoneFunc, error) {
|
|
return func(_ ratelimit.DoneInfo) {
|
|
r.reached = true
|
|
}, nil
|
|
}
|
|
|
|
func (r *ratelimitReachedMock) Allow() (ratelimit.DoneFunc, error) {
|
|
return func(_ ratelimit.DoneInfo) {
|
|
r.reached = true
|
|
}, errors.New("errored")
|
|
}
|
|
|
|
func Test_WithLimiter(t *testing.T) {
|
|
o := options{
|
|
limiter: &ratelimitMock{},
|
|
}
|
|
|
|
WithLimiter(nil)(&o)
|
|
if o.limiter != nil {
|
|
t.Error("The limiter property must be updated.")
|
|
}
|
|
}
|
|
|
|
func Test_Server(t *testing.T) {
|
|
nextValid := func(ctx context.Context, req interface{}) (interface{}, error) {
|
|
return "Hello valid", nil
|
|
}
|
|
|
|
rlm := &ratelimitMock{}
|
|
rlrm := &ratelimitReachedMock{}
|
|
|
|
_, _ = Server(func(o *options) {
|
|
o.limiter = rlm
|
|
})(nextValid)(context.Background(), nil)
|
|
if !rlm.reached {
|
|
t.Error("The ratelimit must run the done function.")
|
|
}
|
|
|
|
_, _ = Server(func(o *options) {
|
|
o.limiter = rlrm
|
|
})(nextValid)(context.Background(), nil)
|
|
if rlrm.reached {
|
|
t.Error("The ratelimit must not run the done function and should be denied.")
|
|
}
|
|
}
|
|
|