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.
131 lines
3.0 KiB
131 lines
3.0 KiB
package rate |
|
|
|
import ( |
|
"context" |
|
"io/ioutil" |
|
"net/http" |
|
"testing" |
|
"time" |
|
|
|
bm "go-common/library/net/http/blademaster" |
|
) |
|
|
|
func TestLimiterUrl(t *testing.T) { |
|
l := New(&Config{ |
|
URLs: map[string]*Limit{"/limit/test": &Limit{Limit: 1, Burst: 2}}, |
|
}) |
|
if !l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should pass,but blocked") |
|
t.FailNow() |
|
} |
|
if !l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should pass,but blocked") |
|
t.FailNow() |
|
} |
|
if l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should block,but passed") |
|
t.FailNow() |
|
} |
|
} |
|
|
|
func TestLimiterApp(t *testing.T) { |
|
l := New(&Config{ |
|
Apps: map[string]*Limit{"testApp": &Limit{Limit: 1, Burst: 2}}, |
|
}) |
|
if !l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should pass,but blocked") |
|
t.FailNow() |
|
} |
|
if !l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should pass,but blocked") |
|
t.FailNow() |
|
} |
|
if l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should block,but passed") |
|
t.FailNow() |
|
} |
|
} |
|
|
|
func TestLimiterUrlApp(t *testing.T) { |
|
l := New(&Config{ |
|
Apps: map[string]*Limit{"testApp": &Limit{Limit: 2, Burst: 1}}, |
|
URLs: map[string]*Limit{"/limit/test": &Limit{Limit: 2, Burst: 1}}, |
|
}) |
|
if !l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should pass,but blocked") |
|
t.FailNow() |
|
} |
|
if l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should block,but passed") |
|
t.FailNow() |
|
} |
|
l.Reload(&Config{ |
|
Apps: map[string]*Limit{"testApp": &Limit{Limit: 1, Burst: 2}}, |
|
URLs: map[string]*Limit{"/limit/test": &Limit{Limit: 1, Burst: 2}}, |
|
}) |
|
if !l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should pass,but blocked") |
|
t.FailNow() |
|
} |
|
if !l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should pass,but blocked") |
|
t.FailNow() |
|
} |
|
if l.Allow("testApp", "/limit/test") { |
|
t.Logf("request should block,but passed") |
|
t.FailNow() |
|
} |
|
} |
|
|
|
func TestLimiterHandler(t *testing.T) { |
|
l := New(&Config{ |
|
Apps: map[string]*Limit{"testApp": &Limit{Limit: 1, Burst: 1}}, |
|
URLs: map[string]*Limit{"/limit/test": &Limit{Limit: 2, Burst: 4}}, |
|
}) |
|
engine := bm.New() |
|
engine.Use(l.Handler()) |
|
engine.GET("/limit/test", func(c *bm.Context) { |
|
c.String(200, "pass") |
|
}) |
|
go engine.Run(":18080") |
|
defer func() { |
|
engine.Server().Shutdown(context.TODO()) |
|
}() |
|
|
|
time.Sleep(time.Millisecond * 20) |
|
code, content, err := httpGet("http://127.0.0.1:18080/limit/test?appkey=testApp") |
|
if err != nil { |
|
t.Logf("http get failed,err:=%v", err) |
|
t.FailNow() |
|
} |
|
if code != 200 || string(content) != "pass" { |
|
t.Logf("request should pass by limiter,but blocked") |
|
t.FailNow() |
|
} |
|
|
|
_, content, err = httpGet("http://127.0.0.1:18080/limit/test?appkey=testApp") |
|
if err != nil { |
|
t.Logf("http get failed,err:=%v", err) |
|
t.FailNow() |
|
} |
|
if string(content) == "pass" { |
|
t.Logf("request should block by limiter,but passed") |
|
t.FailNow() |
|
} |
|
} |
|
|
|
func httpGet(url string) (code int, content []byte, err error) { |
|
resp, err := http.Get(url) |
|
if err != nil { |
|
|
|
return |
|
} |
|
defer resp.Body.Close() |
|
|
|
content, err = ioutil.ReadAll(resp.Body) |
|
if err != nil { |
|
return |
|
} |
|
code = resp.StatusCode |
|
return |
|
}
|
|
|