go-xmpp/tcp_server_mock.go

84 lines
2.1 KiB
Go
Raw Normal View History

package xmpp
2017-10-06 00:11:28 +00:00
import (
"net"
"testing"
)
//=============================================================================
// TCP Server Mock
// ClientHandler is passed by the test client to provide custom behaviour to
// the TCP server mock. This allows customizing the server behaviour to allow
// testing clients under various scenarii.
type ClientHandler func(t *testing.T, conn net.Conn)
// ServerMock is a simple TCP server that can be use to mock basic server
// behaviour to test clients.
type ServerMock struct {
t *testing.T
handler ClientHandler
listener net.Listener
connections []net.Conn
done chan struct{}
}
// Start launches the mock TCP server, listening to an actual address / port.
func (mock *ServerMock) Start(t *testing.T, addr string, handler ClientHandler) {
mock.t = t
mock.handler = handler
if err := mock.init(addr); err != nil {
return
}
go mock.loop()
}
func (mock *ServerMock) Stop() {
close(mock.done)
if mock.listener != nil {
mock.listener.Close()
}
// Close all existing connections
for _, c := range mock.connections {
c.Close()
}
}
//=============================================================================
// Mock Server internals
// init starts listener on the provided address.
func (mock *ServerMock) init(addr string) error {
mock.done = make(chan struct{})
l, err := net.Listen("tcp", addr)
if err != nil {
mock.t.Errorf("TCPServerMock cannot listen on address: %q", addr)
return err
}
mock.listener = l
return nil
}
// loop accepts connections and creates a go routine per connection.
// The go routine is running the client handler, that is used to provide the
// real TCP server behaviour.
func (mock *ServerMock) loop() {
listener := mock.listener
for {
conn, err := listener.Accept()
if err != nil {
select {
case <-mock.done:
return
default:
mock.t.Error("TCPServerMock accept error:", err.Error())
}
return
}
mock.connections = append(mock.connections, conn)
// TODO Create and pass a context to cancel the handler if they are still around = avoid possible leak on complex handlers
go mock.handler(mock.t, conn)
}
}