github.com/xmidt-org/webpa-common@v1.11.9/conlimiter/conlimiter_test.go (about)

     1  package conlimiter
     2  
     3  import (
     4  	//"fmt"
     5  	"github.com/stretchr/testify/assert"
     6  	"github.com/stretchr/testify/mock"
     7  	"net"
     8  	"net/http"
     9  	"sync"
    10  	"sync/atomic"
    11  	"testing"
    12  	"time"
    13  )
    14  
    15  type MyMockedConnection struct {
    16  	mock.Mock
    17  	CloseCalled int32
    18  }
    19  
    20  func (m *MyMockedConnection) Close() error {
    21  	atomic.AddInt32(&m.CloseCalled, 1)
    22  	return nil
    23  }
    24  
    25  func (m *MyMockedConnection) Read(b []byte) (int, error) {
    26  	return 0, nil
    27  }
    28  
    29  func (m *MyMockedConnection) Write(b []byte) (int, error) {
    30  	return 0, nil
    31  }
    32  
    33  func (m *MyMockedConnection) LocalAddr() net.Addr {
    34  	return nil
    35  }
    36  
    37  func (m *MyMockedConnection) RemoteAddr() net.Addr {
    38  	return nil
    39  }
    40  
    41  func (m *MyMockedConnection) SetDeadline(t time.Time) error {
    42  	return nil
    43  }
    44  
    45  func (m *MyMockedConnection) SetReadDeadline(t time.Time) error {
    46  	return nil
    47  }
    48  
    49  func (m *MyMockedConnection) SetWriteDeadline(t time.Time) error {
    50  	return nil
    51  }
    52  
    53  func TestLimit(t *testing.T) {
    54  	testObj := new(MyMockedConnection)
    55  	assert := assert.New(t)
    56  
    57  	s := &http.Server{Addr: ":61234"}
    58  	cl := &ConLimiter{Max: 10}
    59  	cl.Limit(s)
    60  
    61  	for i := 0; i < 10; i++ {
    62  		(s.ConnState)(testObj, http.StateNew)
    63  		assert.Equal(int32(0), atomic.LoadInt32(&testObj.CloseCalled), "Expecting no call.")
    64  	}
    65  	for i := 0; i < 10; i++ {
    66  		(s.ConnState)(testObj, http.StateNew)
    67  		assert.Equal(int32(1), atomic.LoadInt32(&testObj.CloseCalled), "Expecting 1 call.")
    68  		atomic.StoreInt32(&testObj.CloseCalled, 0)
    69  		(s.ConnState)(testObj, http.StateClosed)
    70  		assert.Equal(int32(0), atomic.LoadInt32(&testObj.CloseCalled), "Expecting no call.")
    71  	}
    72  	for i := 0; i < 10; i++ {
    73  		(s.ConnState)(testObj, http.StateHijacked)
    74  		assert.Equal(int32(0), atomic.LoadInt32(&testObj.CloseCalled), "Expecting no call.")
    75  	}
    76  
    77  	assert.Equal(int32(0), atomic.LoadInt32(&testObj.CloseCalled), "Expecting it to be 0.")
    78  	wg := &sync.WaitGroup{}
    79  	for i := 0; i < 100; i++ {
    80  		wg.Add(3)
    81  		go (func() {
    82  			defer wg.Done()
    83  			for i := 0; i < 1000; i++ {
    84  				(s.ConnState)(testObj, http.StateNew)
    85  				(s.ConnState)(testObj, http.StateClosed)
    86  			}
    87  		})()
    88  		go (func() {
    89  			defer wg.Done()
    90  			for i := 0; i < 1000; i++ {
    91  				(s.ConnState)(testObj, http.StateNew)
    92  				(s.ConnState)(testObj, http.StateActive)
    93  				(s.ConnState)(testObj, http.StateClosed)
    94  			}
    95  		})()
    96  		go (func() {
    97  			defer wg.Done()
    98  			for i := 0; i < 1000; i++ {
    99  				(s.ConnState)(testObj, http.StateNew)
   100  				(s.ConnState)(testObj, http.StateActive)
   101  				(s.ConnState)(testObj, http.StateHijacked)
   102  			}
   103  		})()
   104  	}
   105  	wg.Wait()
   106  	assert.Equal(int32(0), cl.current, "Expecting it to be 0.")
   107  }