Add TCP Middlewares support

pkg/tcp/chain_test.go

@@ -0,0 +1,176 @@
+package tcp
+
+import (
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type HandlerTCPFunc func(WriteCloser)
+
+// ServeTCP calls f(conn).
+func (f HandlerTCPFunc) ServeTCP(conn WriteCloser) {
+	f(conn)
+}
+
+// A constructor for middleware
+// that writes its own "tag" into the Conn and does nothing else.
+// Useful in checking if a chain is behaving in the right order.
+func tagMiddleware(tag string) Constructor {
+	return func(h Handler) (Handler, error) {
+		return HandlerTCPFunc(func(conn WriteCloser) {
+			_, err := conn.Write([]byte(tag))
+			if err != nil {
+				panic("Unexpected")
+			}
+			h.ServeTCP(conn)
+		}), nil
+	}
+}
+
+var testApp = HandlerTCPFunc(func(conn WriteCloser) {
+	_, err := conn.Write([]byte("app\n"))
+	if err != nil {
+		panic("unexpected")
+	}
+})
+
+type myWriter struct {
+	data []byte
+}
+
+func (mw *myWriter) Close() error {
+	panic("implement me")
+}
+
+func (mw *myWriter) LocalAddr() net.Addr {
+	panic("implement me")
+}
+
+func (mw *myWriter) RemoteAddr() net.Addr {
+	panic("implement me")
+}
+
+func (mw *myWriter) SetDeadline(t time.Time) error {
+	panic("implement me")
+}
+
+func (mw *myWriter) SetReadDeadline(t time.Time) error {
+	panic("implement me")
+}
+
+func (mw *myWriter) SetWriteDeadline(t time.Time) error {
+	panic("implement me")
+}
+
+func (mw *myWriter) Read(b []byte) (n int, err error) {
+	panic("implement me")
+}
+
+func (mw *myWriter) Write(b []byte) (n int, err error) {
+	mw.data = append(mw.data, b...)
+	return len(mw.data), nil
+}
+
+func (mw *myWriter) CloseWrite() error {
+	return nil
+}
+
+func TestNewChain(t *testing.T) {
+	c1 := func(h Handler) (Handler, error) {
+		return nil, nil
+	}
+
+	c2 := func(h Handler) (Handler, error) {
+		return h, nil
+	}
+
+	slice := []Constructor{c1, c2}
+
+	chain := NewChain(slice...)
+	for k := range slice {
+		assert.ObjectsAreEqual(chain.constructors[k], slice[k])
+	}
+}
+
+func TestThenWorksWithNoMiddleware(t *testing.T) {
+	handler, err := NewChain().Then(testApp)
+	require.NoError(t, err)
+
+	assert.ObjectsAreEqual(handler, testApp)
+}
+
+func TestThenTreatsNilAsError(t *testing.T) {
+	handler, err := NewChain().Then(nil)
+	require.Error(t, err)
+	assert.Nil(t, handler)
+}
+
+func TestThenOrdersHandlersCorrectly(t *testing.T) {
+	t1 := tagMiddleware("t1\n")
+	t2 := tagMiddleware("t2\n")
+	t3 := tagMiddleware("t3\n")
+
+	chained, err := NewChain(t1, t2, t3).Then(testApp)
+	require.NoError(t, err)
+
+	conn := &myWriter{}
+	chained.ServeTCP(conn)
+
+	assert.Equal(t, "t1\nt2\nt3\napp\n", string(conn.data))
+}
+
+func TestAppendAddsHandlersCorrectly(t *testing.T) {
+	chain := NewChain(tagMiddleware("t1\n"), tagMiddleware("t2\n"))
+	newChain := chain.Append(tagMiddleware("t3\n"), tagMiddleware("t4\n"))
+
+	assert.Len(t, chain.constructors, 2)
+	assert.Len(t, newChain.constructors, 4)
+
+	chained, err := newChain.Then(testApp)
+	require.NoError(t, err)
+
+	conn := &myWriter{}
+	chained.ServeTCP(conn)
+
+	assert.Equal(t, "t1\nt2\nt3\nt4\napp\n", string(conn.data))
+}
+
+func TestAppendRespectsImmutability(t *testing.T) {
+	chain := NewChain(tagMiddleware(""))
+	newChain := chain.Append(tagMiddleware(""))
+
+	if &chain.constructors[0] == &newChain.constructors[0] {
+		t.Error("Apppend does not respect immutability")
+	}
+}
+
+func TestExtendAddsHandlersCorrectly(t *testing.T) {
+	chain1 := NewChain(tagMiddleware("t1\n"), tagMiddleware("t2\n"))
+	chain2 := NewChain(tagMiddleware("t3\n"), tagMiddleware("t4\n"))
+	newChain := chain1.Extend(chain2)
+
+	assert.Len(t, chain1.constructors, 2)
+	assert.Len(t, chain2.constructors, 2)
+	assert.Len(t, newChain.constructors, 4)
+
+	chained, err := newChain.Then(testApp)
+	require.NoError(t, err)
+
+	conn := &myWriter{}
+	chained.ServeTCP(conn)
+
+	assert.Equal(t, "t1\nt2\nt3\nt4\napp\n", string(conn.data))
+}
+
+func TestExtendRespectsImmutability(t *testing.T) {
+	chain := NewChain(tagMiddleware(""))
+	newChain := chain.Extend(NewChain(tagMiddleware("")))
+
+	if &chain.constructors[0] == &newChain.constructors[0] {
+		t.Error("Extend does not respect immutability")
+	}
+}