Simulate a tcp connection in Go

2019-03-16 16:23发布

In Go, a TCP connection (net.Conn) is a io.ReadWriteCloser. I'd like to test my network code by simulating a TCP connection. There are two requirements that I have:

  1. the data to be read is stored in a string
  2. whenever data is written, I'd like it to be stored in some kind of buffer which I can access later

Is there a data structure for this, or an easy way to make one?

3条回答
男人必须洒脱
2楼-- · 2019-03-16 17:12

EDIT: I've rolled this answer into a package which makes things a bit simpler - see here: https://github.com/jordwest/mock-conn


While Ivan's solution will work for simple cases, keep in mind that a real TCP connection is actually two buffers, or rather pipes. For example:

 Server   |   Client
 ---------+---------
  reads <===  writes
 writes ===>  reads

If you use a single buffer that the server both reads from and writes to, you could end up with the server talking to itself.

Here's a solution that allows you to pass a MockConn type as a ReadWriteCloser to the server. The Read, Write and Close functions simply proxy through to the functions on the server's end of the pipes.

type MockConn struct {
    ServerReader *io.PipeReader
    ServerWriter *io.PipeWriter

    ClientReader *io.PipeReader
    ClientWriter *io.PipeWriter
}

func (c MockConn) Close() error {
    if err := c.ServerWriter.Close(); err != nil {
        return err
    }
    if err := c.ServerReader.Close(); err != nil {
        return err
    }
    return nil
}

func (c MockConn) Read(data []byte) (n int, err error)  { return c.ServerReader.Read(data) }
func (c MockConn) Write(data []byte) (n int, err error) { return c.ServerWriter.Write(data) }

func NewMockConn() MockConn {
    serverRead, clientWrite := io.Pipe()
    clientRead, serverWrite := io.Pipe()

    return MockConn{
        ServerReader: serverRead,
        ServerWriter: serverWrite,
        ClientReader: clientRead,
        ClientWriter: clientWrite,
    }
}

When mocking a 'server' connection, simply pass the MockConn in place of where you would use the net.Conn (this obviously implements the ReadWriteCloser interface only, you could easily add dummy methods for LocalAddr() etc if you need to support the full net.Conn interface)

In your tests you can act as the client by reading and writing to the ClientReader and ClientWriter fields as needed:

func TestTalkToServer(t *testing.T) {
    /*
     * Assumes that NewMockConn has already been called and
     * the server is waiting for incoming data
     */

    // Send a message to the server
    fmt.Fprintf(mockConn.ClientWriter, "Hello from client!\n")

    // Wait for the response from the server
    rd := bufio.NewReader(mockConn.ClientReader)
    line, err := rd.ReadString('\n')

    if line != "Hello from server!" {
        t.Errorf("Server response not as expected: %s\n", line)
    }
}
查看更多
祖国的老花朵
3楼-- · 2019-03-16 17:17

No idea if this existed when the question was asked, but you probably want net.Pipe() which provides you with two full duplex net.Conn instances which are linked to each other

查看更多
冷血范
4楼-- · 2019-03-16 17:22

Why not using bytes.Buffer? It's an io.ReadWriter and has a String method to get the stored data. If you need to make it an io.ReadWriteCloser, you could define you own type:

type CloseableBuffer struct {
    bytes.Buffer
}

and define a Close method:

func (b *CloseableBuffer) Close() error {
    return nil
}
查看更多
登录 后发表回答