traefik/pkg/tcp/chain_test.go

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")
	}
}