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  }