Add tests for circuit breaker

mtglib/network/circuit_breaker.go

@@ -71,7 +71,7 @@ func (c *circuitBreakerDialer) doClosed(ctx context.Context,
 
 	c.failuresCount++
 
-	if c.state == circuitBreakerStateClosed && c.failuresCount > c.openThreshold {
+	if c.state == circuitBreakerStateClosed && c.failuresCount >= c.openThreshold {
 		c.switchState(circuitBreakerStateOpened)
 	}
 
@@ -184,6 +184,7 @@ func newCircuitBreakerDialer(baseDialer Dialer,
 	openThreshold uint32, halfOpenTimeout, resetFailuresTimeout time.Duration) Dialer {
 	cb := &circuitBreakerDialer{
 		Dialer:               baseDialer,
+		stateMutexChan:       make(chan bool, 1),
 		openThreshold:        openThreshold,
 		halfOpenTimeout:      halfOpenTimeout,
 		resetFailuresTimeout: resetFailuresTimeout,

mtglib/network/circuit_breaker_internal_test.go

@@ -0,0 +1,139 @@
+package network
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/suite"
+)
+
+type CircuitBreakerTestSuite struct {
+	suite.Suite
+
+	d              Dialer
+	mutex          sync.Mutex
+	ctx            context.Context
+	ctxCancel      context.CancelFunc
+	connMock       *ConnMock
+	baseDialerMock *DialerMock
+}
+
+func (suite *CircuitBreakerTestSuite) SetupTest() {
+	suite.mutex = sync.Mutex{}
+	suite.ctx, suite.ctxCancel = context.WithCancel(context.Background())
+	suite.baseDialerMock = &DialerMock{}
+	suite.connMock = &ConnMock{}
+	suite.d = newCircuitBreakerDialer(suite.baseDialerMock,
+		3, 100*time.Millisecond, 50*time.Millisecond)
+}
+
+func (suite *CircuitBreakerTestSuite) TearDownTest() {
+	suite.ctxCancel()
+	suite.baseDialerMock.AssertExpectations(suite.T())
+	suite.connMock.AssertExpectations(suite.T())
+}
+
+func (suite *CircuitBreakerTestSuite) TestMultipleRunsOk() {
+	suite.connMock.On("RemoteAddr").
+		Times(5).
+		Return(&net.TCPAddr{
+			IP:   net.ParseIP("127.0.0.1"),
+			Port: 3128,
+		})
+	suite.baseDialerMock.On("DialContext", mock.Anything, "tcp", "127.0.0.1").
+		Times(5).
+		Return(suite.connMock, nil)
+
+	wg := &sync.WaitGroup{}
+
+	wg.Add(5)
+
+	go func() {
+		wg.Wait()
+		suite.ctxCancel()
+	}()
+
+	for i := 0; i < 5; i++ {
+		go func() {
+			defer wg.Done()
+
+			conn, err := suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+
+			suite.mutex.Lock()
+			defer suite.mutex.Unlock()
+
+			suite.NoError(err)
+			suite.Equal("127.0.0.1:3128", conn.RemoteAddr().String())
+		}()
+	}
+
+	suite.Eventually(func() bool {
+		_, ok := <-suite.ctx.Done()
+
+		return !ok
+	}, time.Second, 10*time.Millisecond)
+}
+
+func (suite *CircuitBreakerTestSuite) TestFromClosedToOpen() {
+	suite.baseDialerMock.On("DialContext", mock.Anything, "tcp", "127.0.0.1").
+		Times(3).
+		Return(&net.TCPConn{}, io.EOF)
+
+	_, err := suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+	suite.True(errors.Is(err, io.EOF))
+
+	_, err = suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+	suite.True(errors.Is(err, io.EOF))
+
+	_, err = suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+	suite.True(errors.Is(err, io.EOF))
+
+	_, err = suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+	suite.True(errors.Is(err, ErrCircuitBreakerOpened))
+}
+
+func (suite *CircuitBreakerTestSuite) TestHalfOpen() {
+	suite.baseDialerMock.On("DialContext", mock.Anything, "tcp", "127.0.0.1").
+		Times(4).
+		Return(&net.TCPConn{}, io.EOF)
+	suite.baseDialerMock.On("DialContext", mock.Anything, "tcp", "127.0.0.2").
+		Twice().
+		Return(suite.connMock, nil)
+	suite.connMock.On("RemoteAddr").Return(&net.TCPAddr{
+		IP:   net.ParseIP("10.0.0.10"),
+		Port: 80,
+	})
+
+	suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+	suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+	suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+	suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+
+	time.Sleep(500 * time.Millisecond)
+
+	_, err := suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+	suite.True(errors.Is(err, io.EOF))
+
+	_, err = suite.d.DialContext(suite.ctx, "tcp", "127.0.0.1")
+	suite.True(errors.Is(err, ErrCircuitBreakerOpened))
+
+	time.Sleep(500 * time.Millisecond)
+
+	conn, err := suite.d.DialContext(suite.ctx, "tcp", "127.0.0.2")
+
+	suite.NoError(err)
+	suite.Equal("10.0.0.10:80", conn.RemoteAddr().String())
+
+	_, err = suite.d.DialContext(suite.ctx, "tcp", "127.0.0.2")
+	suite.NoError(err)
+}
+
+func TestCircuitBreaker(t *testing.T) {
+	suite.Run(t, &CircuitBreakerTestSuite{})
+}

mtglib/network/init_internal_test.go

@@ -0,0 +1,65 @@
+package network
+
+import (
+	"context"
+	"net"
+	"time"
+
+	"github.com/stretchr/testify/mock"
+)
+
+type ConnMock struct {
+	mock.Mock
+}
+
+func (c *ConnMock) Read(b []byte) (int, error) {
+	args := c.Called(b)
+
+	return args.Int(0), args.Error(1)
+}
+
+func (c *ConnMock) Write(b []byte) (int, error) {
+	args := c.Called(b)
+
+	return args.Int(0), args.Error(1)
+}
+
+func (c *ConnMock) Close() error {
+	return c.Called().Error(0)
+}
+
+func (c *ConnMock) LocalAddr() net.Addr {
+	return c.Called().Get(0).(net.Addr)
+}
+
+func (c *ConnMock) RemoteAddr() net.Addr {
+	return c.Called().Get(0).(net.Addr)
+}
+
+func (c *ConnMock) SetDeadline(t time.Time) error {
+	return c.Called(t).Error(0)
+}
+
+func (c *ConnMock) SetReadDeadline(t time.Time) error {
+	return c.Called(t).Error(0)
+}
+
+func (c *ConnMock) SetWriteDeadline(t time.Time) error {
+	return c.Called(t).Error(0)
+}
+
+type DialerMock struct {
+	mock.Mock
+}
+
+func (d *DialerMock) Dial(network, address string) (net.Conn, error) {
+	args := d.Called(network, address)
+
+	return args.Get(0).(net.Conn), args.Error(1)
+}
+
+func (d *DialerMock) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	args := d.Called(ctx, network, address)
+
+	return args.Get(0).(net.Conn), args.Error(1)
+}

mtglib/network/init_test.go

@@ -1,73 +1,13 @@
 package network_test
 
 import (
-	"context"
-	"net"
 	"net/http/httptest"
 	"strings"
-	"time"
 
 	"github.com/mccutchen/go-httpbin/httpbin"
-	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 )
 
-type ConnMock struct {
-	mock.Mock
-}
-
-func (c *ConnMock) Read(b []byte) (int, error) {
-	args := c.Called(b)
-
-	return args.Int(0), args.Error(1)
-}
-
-func (c *ConnMock) Write(b []byte) (int, error) {
-	args := c.Called(b)
-
-	return args.Int(0), args.Error(1)
-}
-
-func (c *ConnMock) Close() error {
-	return c.Called().Error(0)
-}
-
-func (c *ConnMock) LocalAddr() net.Addr {
-	return c.Called().Get(0).(net.Addr)
-}
-
-func (c *ConnMock) RemoteAddr() net.Addr {
-	return c.Called().Get(0).(net.Addr)
-}
-
-func (c *ConnMock) SetDeadline(t time.Time) error {
-	return c.Called(t).Error(0)
-}
-
-func (c *ConnMock) SetReadDeadline(t time.Time) error {
-	return c.Called(t).Error(0)
-}
-
-func (c *ConnMock) SetWriteDeadline(t time.Time) error {
-	return c.Called(t).Error(0)
-}
-
-type DialerMock struct {
-	mock.Mock
-}
-
-func (d *DialerMock) Dial(network, address string) (net.Conn, error) {
-	args := d.Called(network, address)
-
-	return args.Get(0).(net.Conn), args.Error(1)
-}
-
-func (d *DialerMock) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
-	args := d.Called(ctx, network, address)
-
-	return args.Get(0).(net.Conn), args.Error(1)
-}
-
 type HTTPServerTestSuite struct {
 	suite.Suite