github.com/angenalZZZ/gofunc@v0.0.0-20210507121333-48ff1be3917b/http/fast/middleware/limiter/ratelimit_test.go (about)

     1  package limiter
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/angenalZZZ/gofunc/f"
     6  	"github.com/angenalZZZ/gofunc/http/fast"
     7  	"net/http"
     8  	"testing"
     9  	"time"
    10  )
    11  
    12  func TestRateLimit(t *testing.T) {
    13  	// Request times
    14  	times, n := 10000, 0
    15  	rq := f.TimesRepeat(times, http.StatusOK)
    16  	f.TimesRepeatAppend(rq, 100, http.StatusTooManyRequests)
    17  
    18  	app := fast.New()
    19  	rl := NewRateLimiterPerSecond(times)
    20  	app.Get("/", rl.Wrap(func(c *fast.Ctx) {
    21  		c.SendString(c.C.URI().String())
    22  	}))
    23  
    24  	// start test
    25  	for _, want := range rq {
    26  		n++
    27  		req, _ := http.NewRequest("GET", fmt.Sprintf("http://google.com?x=%d", n), nil)
    28  		req.Header.Set("x", f.ToString(n))
    29  		res, _ := app.Test(req)
    30  		if want != res.StatusCode {
    31  			t.Error(ErrTooManyRequests)
    32  		}
    33  	}
    34  
    35  	// wait the interval, should be good for more
    36  	time.Sleep(time.Second)
    37  	for _, want := range rq {
    38  		n++
    39  		req, _ := http.NewRequest("GET", fmt.Sprintf("http://google.com?x=%d", n), nil)
    40  		req.Header.Set("x", f.ToString(n))
    41  		res, _ := app.Test(req)
    42  		if want != res.StatusCode {
    43  			t.Error(ErrTooManyRequests)
    44  		}
    45  	}
    46  
    47  	// many concurrent requests
    48  	rq = f.TimesRepeat(2000, http.StatusOK)
    49  	time.Sleep(time.Second)
    50  	for _, want := range rq {
    51  		n++
    52  		go func(n int, want int) {
    53  			req, _ := http.NewRequest("GET", fmt.Sprintf("http://google.com?x=%d", n), nil)
    54  			req.Header.Set("x", f.ToString(n))
    55  			res, _ := app.Test(req)
    56  			if res == nil || want != res.StatusCode {
    57  				t.Error(ErrTooManyRequests)
    58  			}
    59  		}(n, want.(int))
    60  	}
    61  }
    62  
    63  func TestRateLimiterMiddleware(t *testing.T) {
    64  	// Request times
    65  	times, n := 100000, 0
    66  	rq := f.TimesRepeat(times, http.StatusOK)
    67  	f.TimesRepeatAppend(rq, 100, http.StatusTooManyRequests)
    68  
    69  	app := fast.New()
    70  	app.Use(New(Config{
    71  		Max: int64(times),
    72  		Handler: func(c *fast.Ctx) {
    73  			c.Status(http.StatusTooManyRequests).SendString(c.C.URI().String())
    74  		},
    75  	}))
    76  	app.Get("/", func(c *fast.Ctx) {
    77  		c.SendString(c.C.URI().String())
    78  	})
    79  
    80  	// start test
    81  	for _, want := range rq {
    82  		n++
    83  		req, _ := http.NewRequest("GET", fmt.Sprintf("http://google.com?x=%d", n), nil)
    84  		req.Header.Set("x", f.ToString(n))
    85  		res, _ := app.Test(req)
    86  		if want != res.StatusCode {
    87  			t.Error(ErrTooManyRequests)
    88  		}
    89  	}
    90  
    91  	// wait the interval, should be good for more
    92  	time.Sleep(1 * time.Second)
    93  	for _, want := range rq {
    94  		n++
    95  		req, _ := http.NewRequest("GET", fmt.Sprintf("http://google.com?x=%d", n), nil)
    96  		req.Header.Set("x", f.ToString(n))
    97  		res, _ := app.Test(req)
    98  		if want != res.StatusCode {
    99  			t.Error(ErrTooManyRequests)
   100  		}
   101  	}
   102  
   103  	// many concurrent requests
   104  	rq = f.TimesRepeat(10000, http.StatusOK)
   105  	time.Sleep(time.Second)
   106  	for _, want := range rq {
   107  		n++
   108  		go func(n int, want int) {
   109  			req, _ := http.NewRequest("GET", fmt.Sprintf("http://google.com?x=%d", n), nil)
   110  			req.Header.Set("x", f.ToString(n))
   111  			res, _ := app.Test(req)
   112  			if res == nil || want != res.StatusCode {
   113  				t.Error(ErrTooManyRequests)
   114  			}
   115  		}(n, want.(int))
   116  	}
   117  }