github.com/GuanceCloud/cliutils@v1.1.21/network/http/api_wrap_test.go (about)

     1  // Unless explicitly stated otherwise all files in this repository are licensed
     2  // under the MIT License.
     3  // This product includes software developed at Guance Cloud (https://www.guance.com/).
     4  // Copyright 2021-present Guance, Inc.
     5  
     6  package http
     7  
     8  import (
     9  	"fmt"
    10  	"io"
    11  	"net/http"
    12  	"net/http/httptest"
    13  	"sync"
    14  	"testing"
    15  	"time"
    16  
    17  	tu "github.com/GuanceCloud/cliutils/testutil"
    18  	"github.com/gin-gonic/gin"
    19  )
    20  
    21  type apiStat struct {
    22  	total     int
    23  	costTotal time.Duration
    24  }
    25  
    26  func TestHTTPWrapperWithMetricReporter(t *testing.T) {
    27  	r := gin.New()
    28  
    29  	limitRate := 10
    30  	lmt := NewAPIRateLimiter(float64(limitRate), DefaultRequestKey)
    31  
    32  	testHandler := func(http.ResponseWriter, *http.Request, ...interface{}) (interface{}, error) {
    33  		return nil, nil
    34  	}
    35  
    36  	plg := &WrapPlugins{
    37  		Limiter:  lmt,
    38  		Reporter: &ReporterImpl{},
    39  	}
    40  
    41  	wg := sync.WaitGroup{}
    42  	wg.Add(1)
    43  	go func() {
    44  		defer wg.Done()
    45  		StartReporter()
    46  	}()
    47  
    48  	r.GET("/test", HTTPAPIWrapper(plg, testHandler))
    49  
    50  	ts := httptest.NewServer(r)
    51  	defer ts.Close()
    52  
    53  	time.Sleep(time.Second)
    54  
    55  	var resp *http.Response
    56  	var err error
    57  	var body []byte
    58  	for i := 0; i < limitRate*1000; i++ { // this should exceed max limit and got a 429 status code
    59  		resp, err = http.Get(fmt.Sprintf("%s/test?token=12345", ts.URL))
    60  		if err != nil {
    61  			t.Error(err)
    62  		}
    63  
    64  		body, err = io.ReadAll(resp.Body)
    65  		if err != nil {
    66  			t.Error(err)
    67  		}
    68  
    69  		resp.Body.Close()
    70  	}
    71  
    72  	tu.Equals(t, resp.StatusCode, 429)
    73  	t.Logf("%s", string(body))
    74  
    75  	stats := GetStats()
    76  	for k, v := range stats {
    77  		tu.Assert(t, v.Total == limitRate*1000, "expect %d == %d", v.Total, limitRate)
    78  		tu.Assert(t, v.Limited > 0 && v.Status4XX > 0, "expect %d == %d", v.Total, limitRate)
    79  		t.Logf("%s: %+#v", k, v)
    80  	}
    81  
    82  	StopReporter()
    83  	wg.Wait()
    84  }
    85  
    86  func TestHTTPWrapperWithRateLimit(t *testing.T) {
    87  	r := gin.New()
    88  
    89  	limitRate := 10
    90  	lmt := NewAPIRateLimiter(float64(limitRate), DefaultRequestKey)
    91  
    92  	testHandler := func(http.ResponseWriter, *http.Request, ...interface{}) (interface{}, error) {
    93  		return nil, nil
    94  	}
    95  
    96  	r.GET("/test", HTTPAPIWrapper(&WrapPlugins{
    97  		Limiter: lmt,
    98  	}, testHandler))
    99  
   100  	ts := httptest.NewServer(r)
   101  	defer ts.Close()
   102  
   103  	time.Sleep(time.Second)
   104  
   105  	var resp *http.Response
   106  	var err error
   107  	var body []byte
   108  	for i := 0; i < limitRate*2; i++ { // this should exceed max limit and got a 429 status code
   109  		resp, err = http.Get(fmt.Sprintf("%s/test", ts.URL))
   110  		if err != nil {
   111  			t.Error(err)
   112  		}
   113  
   114  		body, err = io.ReadAll(resp.Body)
   115  		if err != nil {
   116  			t.Error(err)
   117  		}
   118  
   119  		resp.Body.Close()
   120  	}
   121  
   122  	tu.Equals(t, resp.StatusCode, 429)
   123  	t.Logf("%s", string(body))
   124  }