179 lines
3.2 KiB
Go
179 lines
3.2 KiB
Go
package mock
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"math/rand"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/micro/go-micro/transport"
|
|
)
|
|
|
|
type mockSocket struct {
|
|
recv chan *transport.Message
|
|
send chan *transport.Message
|
|
// sock exit
|
|
exit chan bool
|
|
// listener exit
|
|
lexit chan bool
|
|
}
|
|
|
|
type mockClient struct {
|
|
*mockSocket
|
|
opts transport.DialOptions
|
|
}
|
|
|
|
type mockListener struct {
|
|
addr string
|
|
exit chan bool
|
|
conn chan *mockSocket
|
|
opts transport.ListenOptions
|
|
}
|
|
|
|
type mockTransport struct {
|
|
opts transport.Options
|
|
|
|
sync.Mutex
|
|
listeners map[string]*mockListener
|
|
}
|
|
|
|
func (ms *mockSocket) Recv(m *transport.Message) error {
|
|
select {
|
|
case <-ms.exit:
|
|
return errors.New("connection closed")
|
|
case <-ms.lexit:
|
|
return errors.New("server connection closed")
|
|
case cm := <-ms.recv:
|
|
*m = *cm
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (ms *mockSocket) Send(m *transport.Message) error {
|
|
select {
|
|
case <-ms.exit:
|
|
return errors.New("connection closed")
|
|
case <-ms.lexit:
|
|
return errors.New("server connection closed")
|
|
case ms.send <- m:
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (ms *mockSocket) Close() error {
|
|
close(ms.exit)
|
|
return nil
|
|
}
|
|
|
|
func (m *mockListener) Addr() string {
|
|
return m.addr
|
|
}
|
|
|
|
func (m *mockListener) Close() error {
|
|
close(m.exit)
|
|
return nil
|
|
}
|
|
|
|
func (m *mockListener) Accept(fn func(transport.Socket)) error {
|
|
for {
|
|
select {
|
|
case <-m.exit:
|
|
return nil
|
|
case c := <-m.conn:
|
|
go fn(&mockSocket{
|
|
lexit: c.lexit,
|
|
exit: c.exit,
|
|
send: c.recv,
|
|
recv: c.send,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func (m *mockTransport) Dial(addr string, opts ...transport.DialOption) (transport.Client, error) {
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
|
|
listener, ok := m.listeners[addr]
|
|
if !ok {
|
|
return nil, errors.New("could not dial " + addr)
|
|
}
|
|
|
|
var options transport.DialOptions
|
|
for _, o := range opts {
|
|
o(&options)
|
|
}
|
|
|
|
client := &mockClient{
|
|
&mockSocket{
|
|
send: make(chan *transport.Message),
|
|
recv: make(chan *transport.Message),
|
|
exit: make(chan bool),
|
|
lexit: listener.exit,
|
|
},
|
|
options,
|
|
}
|
|
|
|
// pseudo connect
|
|
select {
|
|
case <-listener.exit:
|
|
return nil, errors.New("connection error")
|
|
case listener.conn <- client.mockSocket:
|
|
}
|
|
|
|
return client, nil
|
|
}
|
|
|
|
func (m *mockTransport) Listen(addr string, opts ...transport.ListenOption) (transport.Listener, error) {
|
|
m.Lock()
|
|
defer m.Unlock()
|
|
|
|
var options transport.ListenOptions
|
|
for _, o := range opts {
|
|
o(&options)
|
|
}
|
|
|
|
parts := strings.Split(addr, ":")
|
|
|
|
// if zero port then randomly assign one
|
|
if len(parts) > 1 && parts[len(parts)-1] == "0" {
|
|
r := rand.New(rand.NewSource(time.Now().UnixNano()))
|
|
i := r.Intn(10000)
|
|
// set addr with port
|
|
addr = fmt.Sprintf("%s:%d", parts[:len(parts)-1], 10000+i)
|
|
}
|
|
|
|
if _, ok := m.listeners[addr]; ok {
|
|
return nil, errors.New("already listening on " + addr)
|
|
}
|
|
|
|
listener := &mockListener{
|
|
opts: options,
|
|
addr: addr,
|
|
conn: make(chan *mockSocket),
|
|
exit: make(chan bool),
|
|
}
|
|
|
|
m.listeners[addr] = listener
|
|
|
|
return listener, nil
|
|
}
|
|
|
|
func (m *mockTransport) String() string {
|
|
return "mock"
|
|
}
|
|
|
|
func NewTransport(opts ...transport.Option) transport.Transport {
|
|
var options transport.Options
|
|
for _, o := range opts {
|
|
o(&options)
|
|
}
|
|
|
|
return &mockTransport{
|
|
opts: options,
|
|
listeners: make(map[string]*mockListener),
|
|
}
|
|
}
|