Golang模拟TCP连接做单元测试

下面逐一介绍，笔者接触过的测试手段。

Package bufconn

概要

Taming Complexity in Server/Client(s) Testing in a CI Tool Using gRPC bufconn | Dilip Tadepalli
内部实现是基于管道(pipe)。因为管道是半双工，那么只要两端维持两个管道，就可以模拟了全双工通信了。管道的是由pipe struct表示的, 内部维护一个ring buffer作为缓冲区

// DialContext creates an in-memory full-duplex network connection, unblocks Accept by
// providing it the server half of the connection, and returns the client half
// of the connection.  If ctx is Done, returns ctx.Err()
func (l *Listener) DialContext(ctx context.Context) (net.Conn, error) {
	p1, p2 := newPipe(l.sz), newPipe(l.sz)
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-l.done:
		return nil, errClosed
	case l.ch <- &conn{p1, p2}:
		return &conn{p2, p1}, nil
	}
}

Demo

下面是利用bufconn实现的一个echo服务

package tests

import (
	"bytes"
	"net"
	"testing"
	"tools/bufconn"

	log "github.com/sirupsen/logrus"
)

func echoHandle(conn net.Conn) {
	defer conn.Close()
	buf := bytes.Buffer{}

	for {
		piece := make([]byte, 256)

		n, err := conn.Read(piece)

		piece = piece[:n]

		if err != nil {
			break
		}

		buf.Write(piece)

		n, err = conn.Write(piece)
		if err != nil {
			break
		}

		buf.Next(n)

		log.WithField("count", n).Trace("echo several bytes to client")
	}

	log.Debug("connection closed.")
}

func TestEchoBaseOnBufConn(t *testing.T) {
	kMaxBufferSize := 1024 * 1024
	exitCh := make(chan interface{})

	lis := bufconn.Listen(kMaxBufferSize)
	
	// server 
	go func() {
		defer lis.Close()

		for {
			select {
			case <-exitCh:
				return
			default:
				conn, err := lis.Accept()
				if err != nil {
					log.Warn(err)
				}
				go echoHandle(conn)
			}

		}
	}()

	// client
	cli, _ := lis.Dial()

	kText := "hello"

	buf := bytes.Buffer{}

	piece := make([]byte, 256)

	buf.WriteString(kText)

	for buf.Len() != 0 {
		n, err := cli.Write(buf.Bytes())
		if err != nil {
			log.Debug(err)
			t.Fail()
			break
		}
		buf.Next(n)
	}

	n, err := cli.Read(piece)
	if err != nil {
		log.Debug(err)
		t.Fail()
	}

	piece = piece[:n]

	close(exitCh)

	if string(piece) != kText {
		t.Logf("got a broken string from echo service. expected=%v actual=%v\n", kText, buf.String())
		t.Fail()
	}
}

---

net.Pipe

net.Pipe是由标准库提供的，相比于bufconn在某种程度上使用起来更简洁。如果希望测试单个连接，那么用net.Pipe更方便，如果希望测试快速转换到实际的生产环境，用bufconn更方便。

Demo

只需要将TestEchoBaseOnBufConn略作修改。

func TestEchoBaseOnNetPipe(t *testing.T) {
	exitCh := make(chan interface{})

	cli, serv := net.Pipe()

	go echoHandle(serv)

	kText := "hello"

	buf := bytes.Buffer{}

	piece := make([]byte, 256)

	buf.WriteString(kText)

	for buf.Len() != 0 {
		n, err := cli.Write(buf.Bytes())
		if err != nil {
			log.Debug(err)
			t.Fail()
			break
		}
		buf.Next(n)
	}

	n, err := cli.Read(piece)
	if err != nil {
		log.Debug(err)
		t.Fail()
	}

	piece = piece[:n]

	close(exitCh)

	if string(piece) != kText {
		t.Logf("got a broken string from echo service. expected=%v actual=%v\n", kText, buf.String())
		t.Fail()
	}
}

Package httptest

提供针对于Http的测试模拟，官网文档已经挺详细的。

---

Mocking 🕊

Mock functions in Go

参考链接

How does one test net.Conn in unit tests in Golang?