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 }