openbilibili/library/net/http/blademaster/middleware/rate/limit_test.go

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
}