github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/request/tpslimiter_test.go (about)

     1  package request
     2  
     3  import (
     4  	"context"
     5  	"github.com/stretchr/testify/assert"
     6  	"testing"
     7  	"time"
     8  )
     9  
    10  func TestLimit(t *testing.T) {
    11  	a := assert.New(t)
    12  	l := NewTPSLimiter()
    13  	finished := make(chan struct{})
    14  	go func() {
    15  		l.Limit(context.Background(), "token", 1, 1)
    16  		close(finished)
    17  	}()
    18  	select {
    19  	case <-finished:
    20  	case <-time.After(10 * time.Second):
    21  		a.Fail("Limit should be finished instantly.")
    22  	}
    23  
    24  	finished = make(chan struct{})
    25  	go func() {
    26  		l.Limit(context.Background(), "token", 1, 1)
    27  		close(finished)
    28  	}()
    29  	select {
    30  	case <-finished:
    31  	case <-time.After(2 * time.Second):
    32  		a.Fail("Limit should be finished in 1 second.")
    33  	}
    34  
    35  	finished = make(chan struct{})
    36  	go func() {
    37  		l.Limit(context.Background(), "token", 10, 1)
    38  		close(finished)
    39  	}()
    40  	select {
    41  	case <-finished:
    42  	case <-time.After(1 * time.Second):
    43  		a.Fail("Limit should be finished instantly.")
    44  	}
    45  
    46  }