github.com/lingyao2333/mo-zero@v1.4.1/core/syncx/timeoutlimit_test.go (about)

     1  package syncx
     2  
     3  import (
     4  	"sync"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func TestTimeoutLimit(t *testing.T) {
    12  	limit := NewTimeoutLimit(2)
    13  	assert.Nil(t, limit.Borrow(time.Millisecond*200))
    14  	assert.Nil(t, limit.Borrow(time.Millisecond*200))
    15  	var wait1, wait2, wait3 sync.WaitGroup
    16  	wait1.Add(1)
    17  	wait2.Add(1)
    18  	wait3.Add(1)
    19  	go func() {
    20  		wait1.Wait()
    21  		wait2.Done()
    22  		assert.Nil(t, limit.Return())
    23  		wait3.Done()
    24  	}()
    25  	wait1.Done()
    26  	wait2.Wait()
    27  	assert.Nil(t, limit.Borrow(time.Second))
    28  	wait3.Wait()
    29  	assert.Equal(t, ErrTimeout, limit.Borrow(time.Millisecond*100))
    30  	assert.Nil(t, limit.Return())
    31  	assert.Nil(t, limit.Return())
    32  	assert.Equal(t, ErrLimitReturn, limit.Return())
    33  }