Initial commit

lib/testutil/limitwriter.go

@@ -0,0 +1,33 @@
+package testutil
+
+import (
+	"io"
+	"os"
+)
+
+func NewLimitWriter(w io.Writer, limit int) *LimitWriter {
+	return &LimitWriter{
+		w:     w,
+		limit: limit,
+	}
+}
+
+type LimitWriter struct {
+	w       io.Writer
+	limit   int
+	written int
+}
+
+func (lw *LimitWriter) Write(buf []byte) (int, error) {
+	n, err := lw.w.Write(buf)
+	if err != nil {
+		return n, err
+	}
+
+	lw.written += n
+	if lw.written > lw.limit {
+		return n, os.ErrClosed
+	}
+
+	return n, nil
+}

lib/testutil/testconn.go

@@ -0,0 +1,79 @@
+package testutil
+
+import (
+	"net"
+	"sync"
+	"time"
+)
+
+type Network struct {
+	lock sync.Mutex
+	// Current connections.
+	cConn net.Conn
+	sConn net.Conn
+
+	acceptQ chan net.Conn
+}
+
+func NewNetwork() *Network {
+	return &Network{
+		acceptQ: make(chan net.Conn, 1),
+	}
+}
+
+func (n *Network) Dial() net.Conn {
+	cc, sc := net.Pipe()
+	func() {
+		n.lock.Lock()
+		defer n.lock.Unlock()
+		if n.cConn != nil {
+			n.cConn.Close()
+			n.cConn = nil
+		}
+		select {
+		case n.acceptQ <- sc:
+			n.cConn = cc
+		default:
+			cc = nil
+		}
+	}()
+	return cc
+}
+
+func (n *Network) Accept() net.Conn {
+	var sc net.Conn
+	select {
+	case sc = <-n.acceptQ:
+	case <-time.After(time.Second):
+		return nil
+	}
+
+	func() {
+		n.lock.Lock()
+		defer n.lock.Unlock()
+		if n.sConn != nil {
+			n.sConn.Close()
+			n.sConn = nil
+		}
+		n.sConn = sc
+	}()
+	return sc
+}
+
+func (n *Network) CloseClient() {
+	n.lock.Lock()
+	defer n.lock.Unlock()
+	if n.cConn != nil {
+		n.cConn.Close()
+		n.cConn = nil
+	}
+}
+
+func (n *Network) CloseServer() {
+	n.lock.Lock()
+	defer n.lock.Unlock()
+	if n.sConn != nil {
+		n.sConn.Close()
+		n.sConn = nil
+	}
+}

lib/testutil/util.go

@@ -0,0 +1,10 @@
+package testutil
+
+import "testing"
+
+func AssertNotNil(t *testing.T, err error) {
+	t.Helper()
+	if err != nil {
+		t.Fatal(err)
+	}
+}