gopkg.in/rethinkdb/rethinkdb-go.v6@v6.2.2/mocks_test.go (about) 1 package rethinkdb 2 3 import ( 4 "github.com/stretchr/testify/mock" 5 "io" 6 "net" 7 "time" 8 ) 9 10 type connMock struct { 11 mock.Mock 12 done <-chan struct{} 13 doneSet chan struct{} 14 } 15 16 func (m *connMock) setDone(done <-chan struct{}) { 17 m.done = done 18 close(m.doneSet) 19 } 20 21 func (m *connMock) waitDial() { 22 <-m.doneSet 23 } 24 25 func (m *connMock) waitDone() { 26 <-m.done 27 } 28 29 func (m *connMock) onCloseReturn(err error) { 30 closeChan := make(chan struct{}) 31 m.doneSet = make(chan struct{}) 32 // Maybe - Connection can be closed by Close() before Read() occurs when stopReadChan closed 33 m.On("Read", respHeaderLen).Return(nil, 0, io.EOF, nil).Maybe().Run(func(args mock.Arguments) { 34 <-closeChan 35 }) 36 m.On("Close").Return(err).Once().Run(func(args mock.Arguments) { 37 close(closeChan) 38 }) 39 } 40 41 func (m *connMock) Read(b []byte) (n int, err error) { 42 args := m.Called(len(b)) 43 rbuf, ok := args.Get(0).([]byte) 44 if ok { 45 copy(b, rbuf) 46 } 47 timeout := args.Get(3) 48 if timeout != nil { 49 time.Sleep(timeout.(time.Duration)) 50 } 51 return args.Int(1), args.Error(2) 52 } 53 54 func (m *connMock) Write(b []byte) (n int, err error) { 55 args := m.Called(b) 56 timeout := args.Get(2) 57 if timeout != nil { 58 time.Sleep(timeout.(time.Duration)) 59 } 60 return args.Int(0), args.Error(1) 61 } 62 63 func (m *connMock) Close() error { 64 args := m.Called() 65 return args.Error(0) 66 } 67 68 func (m *connMock) LocalAddr() net.Addr { 69 args := m.Called() 70 return args.Get(0).(net.Addr) 71 } 72 73 func (m *connMock) RemoteAddr() net.Addr { 74 args := m.Called() 75 return args.Get(0).(net.Addr) 76 } 77 78 func (m *connMock) SetDeadline(t time.Time) error { 79 args := m.Called() 80 return args.Error(0) 81 } 82 83 func (m *connMock) SetReadDeadline(t time.Time) error { 84 args := m.Called() 85 return args.Error(0) 86 } 87 88 func (m *connMock) SetWriteDeadline(t time.Time) error { 89 args := m.Called() 90 return args.Error(0) 91 }