Yahya/5005-fix-flakey-libp2p-test

network/gossip/libp2p/peerManager.go

@@ -3,6 +3,7 @@ package libp2p
 import (
 	"context"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/hashicorp/go-multierror"
@@ -53,14 +54,21 @@ func NewPeerManager(ctx context.Context, logger zerolog.Logger, idsProvider func
 }
 
 // Start kicks off the ambient periodic connection updates
+// It gets blocking till all connection updates are kicking off
 func (pm *PeerManager) Start() error {
-	go pm.updateLoop()
-	go pm.periodicUpdate()
+	wg := &sync.WaitGroup{}
+	wg.Add(2)
+
+	go pm.updateLoop(wg)
+	go pm.periodicUpdate(wg)
+
+	wg.Wait()
 	return nil
 }
 
 // updateLoop triggers an update peer request when it has been requested
-func (pm *PeerManager) updateLoop() {
+func (pm *PeerManager) updateLoop(wg *sync.WaitGroup) {
+	wg.Done()
 	for {
 		select {
 		case <-pm.ctx.Done():
@@ -72,8 +80,8 @@ func (pm *PeerManager) updateLoop() {
 }
 
 // updateLoop request periodic connection update
-func (pm *PeerManager) periodicUpdate() {
-
+func (pm *PeerManager) periodicUpdate(wg *sync.WaitGroup) {
+	wg.Done()
 	// request initial discovery
 	pm.RequestPeerUpdate()
 

network/gossip/libp2p/peerManager_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/assert"
 	testifymock "github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 
 	"github.com/onflow/flow-go/model/flow"
@@ -29,14 +30,14 @@ func TestPeerManagerTestSuite(t *testing.T) {
 	suite.Run(t, new(PeerManagerTestSuite))
 }
 
-func (ts *PeerManagerTestSuite) SetupTest() {
-	ts.log = ts.log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).With().Caller().Logger()
-	ts.ctx = context.Background()
+func (suite *PeerManagerTestSuite) SetupTest() {
+	suite.log = zerolog.New(os.Stderr).Level(zerolog.ErrorLevel)
+	suite.ctx = context.Background()
 }
 
 // TestUpdatePeers tests that updatePeers calls the connector with the expected list of ids to connect and disconnect
 // from. The tests are cumulative and ordered.
-func (ts *PeerManagerTestSuite) TestUpdatePeers() {
+func (suite *PeerManagerTestSuite) TestUpdatePeers() {
 
 	// create some test ids
 	currentIDs := unittest.IdentityListFixture(10)
@@ -51,59 +52,59 @@ func (ts *PeerManagerTestSuite) TestUpdatePeers() {
 
 	// create the connector mock to check ids requested for connect and disconnect
 	connector := new(mock.Connector)
-	connector.On("ConnectPeers", ts.ctx, testifymock.AnythingOfType("flow.IdentityList")).
+	connector.On("ConnectPeers", suite.ctx, testifymock.AnythingOfType("flow.IdentityList")).
 		Run(func(args testifymock.Arguments) {
 			idArg := args[1].(flow.IdentityList)
-			assertListsEqual(ts.T(), currentIDs, idArg)
+			assertListsEqual(suite.T(), currentIDs, idArg)
 		}).
 		Return(nil)
-	connector.On("DisconnectPeers", ts.ctx, testifymock.AnythingOfType("flow.IdentityList")).
+	connector.On("DisconnectPeers", suite.ctx, testifymock.AnythingOfType("flow.IdentityList")).
 		Run(func(args testifymock.Arguments) {
 			idArg := args[1].(flow.IdentityList)
-			assertListsEqual(ts.T(), extraIDs, idArg)
+			assertListsEqual(suite.T(), extraIDs, idArg)
 			// assert that ids passed to disconnect have no id in common with those passed to connect
-			assertListsDisjoint(ts.T(), currentIDs, extraIDs)
+			assertListsDisjoint(suite.T(), currentIDs, extraIDs)
 		}).
 		Return(nil)
 
 	// create the peer manager (but don't start it)
-	pm := NewPeerManager(ts.ctx, ts.log, idProvider, connector)
+	pm := NewPeerManager(suite.ctx, suite.log, idProvider, connector)
 
 	// very first call to updatepeer
-	ts.Run("updatePeers only connects to all peers the first time", func() {
+	suite.Run("updatePeers only connects to all peers the first time", func() {
 
 		pm.updatePeers()
 
-		connector.AssertNumberOfCalls(ts.T(), "ConnectPeers", 1)
-		connector.AssertNotCalled(ts.T(), "DisconnectPeers")
+		connector.AssertNumberOfCalls(suite.T(), "ConnectPeers", 1)
+		connector.AssertNotCalled(suite.T(), "DisconnectPeers")
 	})
 
 	// a subsequent call to updatepeer should request a connect to existing ids and new ids
-	ts.Run("updatePeers connects to old and new peers", func() {
+	suite.Run("updatePeers connects to old and new peers", func() {
 		// create a new id
 		newIDs := unittest.IdentityListFixture(1)
 		currentIDs = append(currentIDs, newIDs...)
 
 		pm.updatePeers()
 
-		connector.AssertNumberOfCalls(ts.T(), "ConnectPeers", 2)
-		connector.AssertNotCalled(ts.T(), "DisconnectPeers")
+		connector.AssertNumberOfCalls(suite.T(), "ConnectPeers", 2)
+		connector.AssertNotCalled(suite.T(), "DisconnectPeers")
 	})
 
 	// when ids are excluded, they should be requested to be disconnected
-	ts.Run("updatePeers disconnects from extra peers", func() {
+	suite.Run("updatePeers disconnects from extra peers", func() {
 		// delete an id
 		extraIDs = currentIDs.Sample(1)
 		currentIDs = currentIDs.Filter(filter.Not(filter.In(extraIDs)))
 
 		pm.updatePeers()
 
-		connector.AssertNumberOfCalls(ts.T(), "ConnectPeers", 3)
-		connector.AssertNumberOfCalls(ts.T(), "DisconnectPeers", 1)
+		connector.AssertNumberOfCalls(suite.T(), "ConnectPeers", 3)
+		connector.AssertNumberOfCalls(suite.T(), "DisconnectPeers", 1)
 	})
 
 	// addition and deletion of ids should result in appropriate connect and disconnect calls
-	ts.Run("updatePeers connects to new peers and disconnects from extra peers", func() {
+	suite.Run("updatePeers connects to new peers and disconnects from extra peers", func() {
 		// remove a couple of ids
 		extraIDs = currentIDs.Sample(2)
 		currentIDs = currentIDs.Filter(filter.Not(filter.In(extraIDs)))
@@ -114,108 +115,128 @@ func (ts *PeerManagerTestSuite) TestUpdatePeers() {
 
 		pm.updatePeers()
 
-		connector.AssertNumberOfCalls(ts.T(), "ConnectPeers", 4)
-		connector.AssertNumberOfCalls(ts.T(), "DisconnectPeers", 2)
+		connector.AssertNumberOfCalls(suite.T(), "ConnectPeers", 4)
+		connector.AssertNumberOfCalls(suite.T(), "DisconnectPeers", 2)
 	})
 }
 
-// TestPeriodicPeerUpdate tests that the peermanager runs periodically
-func (ts *PeerManagerTestSuite) TestPeriodicPeerUpdate() {
+// TestPeriodicPeerUpdate tests that the peer manager runs periodically
+func (suite *PeerManagerTestSuite) TestPeriodicPeerUpdate() {
 	currentIDs := unittest.IdentityListFixture(10)
 	idProvider := func() (flow.IdentityList, error) {
 		return currentIDs, nil
 	}
 
 	connector := new(mock.Connector)
-	connector.On("ConnectPeers", ts.ctx, testifymock.Anything).Return(nil)
-	connector.On("DisconnectPeers", ts.ctx, testifymock.Anything).Return(nil)
-	pm := NewPeerManager(ts.ctx, ts.log, idProvider, connector)
+	wg := &sync.WaitGroup{} // keeps track of number of calls on `ConnectPeers`
+	count, times := 0, 2    // we expect it to be called twice at least
+	wg.Add(times)
+	connector.On("ConnectPeers", suite.ctx, testifymock.Anything).Run(func(args testifymock.Arguments) {
+		if count < times {
+			count++
+			wg.Done()
+		}
+	}).Return(nil)
+	connector.On("DisconnectPeers", suite.ctx, testifymock.Anything).Return(nil)
+	pm := NewPeerManager(suite.ctx, suite.log, idProvider, connector)
 
 	PeerUpdateInterval = 5 * time.Millisecond
 	err := pm.Start()
-	assert.NoError(ts.T(), err)
-	assert.Eventually(ts.T(), func() bool {
-		return connector.AssertNumberOfCalls(ts.T(), "ConnectPeers", 2)
-	}, 2*PeerUpdateInterval+4*time.Millisecond, 2*PeerUpdateInterval)
+	assert.NoError(suite.T(), err)
+
+	unittest.RequireReturnsBefore(suite.T(), wg.Wait, 2*PeerUpdateInterval,
+		"ConnectPeers is not running on UpdateIntervals")
 }
 
 // TestOnDemandPeerUpdate tests that the a peer update can be requested on demand and in between the periodic runs
-func (ts *PeerManagerTestSuite) TestOnDemandPeerUpdate() {
+func (suite *PeerManagerTestSuite) TestOnDemandPeerUpdate() {
 	currentIDs := unittest.IdentityListFixture(10)
 	idProvider := func() (flow.IdentityList, error) {
 		return currentIDs, nil
 	}
 
-	connector := new(mock.Connector)
-	connector.On("ConnectPeers", ts.ctx, testifymock.Anything).Return(nil)
-	connector.On("DisconnectPeers", ts.ctx, testifymock.Anything).Return(nil)
-	pm := NewPeerManager(ts.ctx, ts.log, idProvider, connector)
-
+	// chooses peer interval rate deliberately long to capture on demand peer update
 	PeerUpdateInterval = time.Hour
-	err := pm.Start()
-	assert.NoError(ts.T(), err)
-
-	// wait for the first periodic update initiated after start to complete
-	assert.Eventually(ts.T(), func() bool {
-		return connector.AssertNumberOfCalls(ts.T(), "ConnectPeers", 1)
-	}, 10*time.Millisecond, 1*time.Millisecond)
 
-	// make a request for peer update
+	// creates mock connector
+	wg := &sync.WaitGroup{} // keeps track of number of calls on `ConnectPeers`
+	times, count := 2, 0    // we expect it to be called twice overall
+	wg.Add(1)               // this accounts for one invocation, the other invocation is subsequent
+	connector := new(mock.Connector)
+	// captures the first periodic update initiated after start to complete
+	connector.On("ConnectPeers", suite.ctx, testifymock.Anything).Run(func(args testifymock.Arguments) {
+		if count < times {
+			count++
+			wg.Done()
+		}
+	}).Return(nil)
+	connector.On("DisconnectPeers", suite.ctx, testifymock.Anything).Return(nil)
+	pm := NewPeerManager(suite.ctx, suite.log, idProvider, connector)
+
+	// starts peer manager and waits for first periodic update
+	require.NoError(suite.T(), pm.Start())
+	unittest.RequireReturnsBefore(suite.T(), wg.Wait, 10*time.Second,
+		"ConnectPeers is not running on startup")
+
+	// makes a request for peer update
+	wg.Add(1) // expects a call to `ConnectPeers` by requesting peer update
 	pm.RequestPeerUpdate()
 
 	// assert that a call to connect to peers is made
-	assert.Eventually(ts.T(), func() bool {
-		return connector.AssertNumberOfCalls(ts.T(), "ConnectPeers", 2)
-	}, 10*time.Millisecond, 1*time.Millisecond)
+	unittest.RequireReturnsBefore(suite.T(), wg.Wait, 1*time.Millisecond,
+		"ConnectPeers is not running on request")
 }
 
 // TestConcurrentOnDemandPeerUpdate tests that concurrent on-demand peer update request never block
-func (ts *PeerManagerTestSuite) TestConcurrentOnDemandPeerUpdate() {
+func (suite *PeerManagerTestSuite) TestConcurrentOnDemandPeerUpdate() {
 	currentIDs := unittest.IdentityListFixture(10)
 	idProvider := func() (flow.IdentityList, error) {
 		return currentIDs, nil
 	}
 
-	ctx, cancel := context.WithCancel(ts.ctx)
+	ctx, cancel := context.WithCancel(suite.ctx)
 	defer cancel()
 
 	connector := new(mock.Connector)
 	// connectPeerGate channel gates the return of the connector
 	connectPeerGate := make(chan time.Time)
 	defer close(connectPeerGate)
-	connector.On("ConnectPeers", ctx, testifymock.Anything).Return(nil).WaitUntil(connectPeerGate)
-	connector.On("DisconnectPeers", ctx, testifymock.Anything).Return(nil)
 
-	pm := NewPeerManager(ctx, ts.log, idProvider, connector)
+	// creates mock connector
+	wg := &sync.WaitGroup{} // keeps track of number of calls on `ConnectPeers`
+	// we expect it to be called twice, i.e.,
+	// one for periodic update and one for the on-demand request
+	count, times := 0, 2
+	wg.Add(times)
+	connector.On("ConnectPeers", ctx, testifymock.Anything).Run(func(args testifymock.Arguments) {
+		if count < times {
+			count++
+			wg.Done()
+		}
+	}).Return(nil).
+		WaitUntil(connectPeerGate) // blocks call for connectPeerGate channel
+	connector.On("DisconnectPeers", ctx, testifymock.Anything).Return(nil)
+	pm := NewPeerManager(ctx, suite.log, idProvider, connector)
 
 	// set the periodic interval to a high value so that periodic runs don't interfere with this test
 	PeerUpdateInterval = time.Hour
 
 	// start the peer manager
 	// this should trigger the first update and which will block on the ConnectPeers to return
 	err := pm.Start()
-	assert.NoError(ts.T(), err)
-
-	// make 10 concurrent request for peer update
-	wg := sync.WaitGroup{}
-	for i := 0; i < 10; i++ {
-		wg.Add(1)
-		go func() {
-			pm.RequestPeerUpdate()
-			wg.Done()
-		}()
-	}
+	assert.NoError(suite.T(), err)
 
-	// assert that none of the request is blocked even if update is blocked
-	unittest.AssertReturnsBefore(ts.T(), wg.Wait, time.Second)
+	// makes 10 concurrent request for peer update
+	unittest.RequireConcurrentCallsReturnBefore(suite.T(), pm.RequestPeerUpdate, 10, time.Second,
+		"concurrent peer update requests could not return on time")
 
-	// allow the first update to finish
+	// allow the first and second updates to finish
+	connectPeerGate <- time.Now()
 	connectPeerGate <- time.Now()
 
-	// assert that only two calls to ConnectPeers were made (one for periodic update and one for the on-demand request)
-	assert.Eventually(ts.T(), func() bool {
-		return connector.AssertNumberOfCalls(ts.T(), "ConnectPeers", 2)
-	}, 10*time.Millisecond, 1*time.Millisecond)
+	// requires two calls to ConnectPeers were made
+	unittest.RequireReturnsBefore(suite.T(), wg.Wait, 10*time.Second,
+		"required invocations on ConnectedPeers did not happen")
 }
 
 // assertListsEqual asserts that two identity list are equal ignoring the order

network/gossip/libp2p/test/echoengine_test.go

@@ -8,9 +8,7 @@ import (
 	"testing"
 	"time"
 
-	golog "github.com/ipfs/go-log"
 	"github.com/rs/zerolog"
-	"github.com/rs/zerolog/log"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
@@ -42,8 +40,8 @@ func TestStubEngineTestSuite(t *testing.T) {
 
 func (s *EchoEngineTestSuite) SetupTest() {
 	const count = 2
-	golog.SetAllLoggers(golog.LevelError)
-	logger := log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).With().Caller().Logger()
+
+	logger := zerolog.New(os.Stderr).Level(zerolog.ErrorLevel)
 	s.ids, s.mws, s.nets = generateIDsMiddlewaresNetworks(s.T(), count, logger, 100, nil, false)
 }
 

network/gossip/libp2p/test/epochtransition_test.go

@@ -8,9 +8,7 @@ import (
 	"testing"
 	"time"
 
-	golog "github.com/ipfs/go-log"
 	"github.com/rs/zerolog"
-	"github.com/rs/zerolog/log"
 	"github.com/stretchr/testify/assert"
 	testifymock "github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -48,8 +46,7 @@ func TestEpochTransitionTestSuite(t *testing.T) {
 func (ts *EpochTransitionTestSuite) SetupTest() {
 	rand.Seed(time.Now().UnixNano())
 	nodeCount := 10
-	golog.SetAllLoggers(golog.LevelError)
-	ts.logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).With().Caller().Logger()
+	ts.logger = zerolog.New(os.Stderr).Level(zerolog.ErrorLevel)
 
 	// create ids
 	ids, mws := generateIDsAndMiddlewares(ts.T(), nodeCount, ts.logger)
@@ -106,7 +103,6 @@ func (ts *EpochTransitionTestSuite) TearDownTest() {
 // TestNewNodeAdded tests that an additional node in the next epoch gets connected to other nodes and can exchange messages
 // in the current epoch
 func (ts *EpochTransitionTestSuite) TestNewNodeAdded() {
-
 	// create the id, middleware and network for a new node
 	ids, mws, nets := generateIDsMiddlewaresNetworks(ts.T(), 1, ts.logger, 100, nil, false)
 	newMiddleware := mws[0]
@@ -143,7 +139,6 @@ func (ts *EpochTransitionTestSuite) TestNewNodeAdded() {
 
 // TestNodeRemoved tests that a node that is removed in the next epoch remains connected for the current epoch
 func (ts *EpochTransitionTestSuite) TestNodeRemoved() {
-
 	// choose a random node to remove
 	removeIndex := rand.Intn(len(ts.ids))
 	removedID := ts.ids[removeIndex]