Moves connectPeers to common test setup file (#5800)

* moves connectPeers

beacon-chain/sync/initial-sync/initial_sync_test.go

@@ -2,25 +2,31 @@ package initialsync
 
 import (
 	"context"
+	"fmt"
 	"io/ioutil"
 	"os"
 	"reflect"
+	"sync"
 	"testing"
 	"time"
 
-	"sync"
-
+	"github.com/ethereum/go-ethereum/p2p/enr"
+	"github.com/libp2p/go-libp2p-core/network"
 	eth "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	mock "github.com/prysmaticlabs/prysm/beacon-chain/blockchain/testing"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
 	"github.com/prysmaticlabs/prysm/beacon-chain/db"
 	dbtest "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
+	"github.com/prysmaticlabs/prysm/beacon-chain/p2p/peers"
 	p2pt "github.com/prysmaticlabs/prysm/beacon-chain/p2p/testing"
 	stateTrie "github.com/prysmaticlabs/prysm/beacon-chain/state"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
+	beaconsync "github.com/prysmaticlabs/prysm/beacon-chain/sync"
 	p2ppb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+	"github.com/prysmaticlabs/prysm/shared/hashutil"
 	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/roughtime"
+	"github.com/prysmaticlabs/prysm/shared/sliceutil"
 	"github.com/sirupsen/logrus"
 )
 
@@ -141,3 +147,94 @@ func TestMakeSequence(t *testing.T) {
 		t.Fatalf("Wanted %v, got %v", want, got)
 	}
 }
+
+// Connect peers with local host. This method sets up peer statuses and the appropriate handlers
+// for each test peer.
+func connectPeers(t *testing.T, host *p2pt.TestP2P, data []*peerData, peerStatus *peers.Status) {
+	const topic = "/eth2/beacon_chain/req/beacon_blocks_by_range/1/ssz"
+
+	for _, d := range data {
+		peer := p2pt.NewTestP2P(t)
+
+		// Copy pointer for callback scope.
+		var datum = d
+
+		peer.SetStreamHandler(topic, func(stream network.Stream) {
+			defer func() {
+				if err := stream.Close(); err != nil {
+					t.Log(err)
+				}
+			}()
+
+			req := &p2ppb.BeaconBlocksByRangeRequest{}
+			if err := peer.Encoding().DecodeWithLength(stream, req); err != nil {
+				t.Error(err)
+			}
+
+			requestedBlocks := makeSequence(req.StartSlot, req.StartSlot+(req.Count*req.Step))
+
+			// Expected failure range
+			if len(sliceutil.IntersectionUint64(datum.failureSlots, requestedBlocks)) > 0 {
+				if _, err := stream.Write([]byte{0x01}); err != nil {
+					t.Error(err)
+				}
+				if _, err := peer.Encoding().EncodeWithLength(stream, "bad"); err != nil {
+					t.Error(err)
+				}
+				return
+			}
+
+			// Determine the correct subset of blocks to return as dictated by the test scenario.
+			blocks := sliceutil.IntersectionUint64(datum.blocks, requestedBlocks)
+
+			ret := make([]*eth.SignedBeaconBlock, 0)
+			for _, slot := range blocks {
+				if (slot-req.StartSlot)%req.Step != 0 {
+					continue
+				}
+				cache.RLock()
+				parentRoot := cache.rootCache[cache.parentSlotCache[slot]]
+				cache.RUnlock()
+				blk := &eth.SignedBeaconBlock{
+					Block: &eth.BeaconBlock{
+						Slot:       slot,
+						ParentRoot: parentRoot[:],
+					},
+				}
+				// If forked peer, give a different parent root.
+				if datum.forkedPeer {
+					newRoot := hashutil.Hash(parentRoot[:])
+					blk.Block.ParentRoot = newRoot[:]
+				}
+				ret = append(ret, blk)
+				currRoot, err := stateutil.BlockRoot(blk.Block)
+				if err != nil {
+					t.Fatal(err)
+				}
+				logrus.Infof("block with slot %d , signing root %#x and parent root %#x", slot, currRoot, parentRoot)
+			}
+
+			if uint64(len(ret)) > req.Count {
+				ret = ret[:req.Count]
+			}
+
+			for i := 0; i < len(ret); i++ {
+				if err := beaconsync.WriteChunk(stream, peer.Encoding(), ret[i]); err != nil {
+					t.Error(err)
+				}
+			}
+		})
+
+		peer.Connect(host)
+
+		peerStatus.Add(new(enr.Record), peer.PeerID(), nil, network.DirOutbound)
+		peerStatus.SetConnectionState(peer.PeerID(), peers.PeerConnected)
+		peerStatus.SetChainState(peer.PeerID(), &p2ppb.Status{
+			ForkDigest:     params.BeaconConfig().GenesisForkVersion,
+			FinalizedRoot:  []byte(fmt.Sprintf("finalized_root %d", datum.finalizedEpoch)),
+			FinalizedEpoch: datum.finalizedEpoch,
+			HeadRoot:       []byte("head_root"),
+			HeadSlot:       datum.headSlot,
+		})
+	}
+}

beacon-chain/sync/initial-sync/round_robin_test.go

@@ -2,25 +2,17 @@ package initialsync
 
 import (
 	"context"
-	"fmt"
 	"testing"
 
-	"github.com/ethereum/go-ethereum/p2p/enr"
 	"github.com/kevinms/leakybucket-go"
-	"github.com/libp2p/go-libp2p-core/network"
 	eth "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	mock "github.com/prysmaticlabs/prysm/beacon-chain/blockchain/testing"
 	dbtest "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
-	"github.com/prysmaticlabs/prysm/beacon-chain/p2p/peers"
 	p2pt "github.com/prysmaticlabs/prysm/beacon-chain/p2p/testing"
 	stateTrie "github.com/prysmaticlabs/prysm/beacon-chain/state"
-	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
-	beaconsync "github.com/prysmaticlabs/prysm/beacon-chain/sync"
 	p2ppb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
-	"github.com/prysmaticlabs/prysm/shared/hashutil"
 	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/sliceutil"
-	"github.com/sirupsen/logrus"
 )
 
 func TestConstants(t *testing.T) {
@@ -292,94 +284,3 @@ func TestRoundRobinSync(t *testing.T) {
 		})
 	}
 }
-
-// Connect peers with local host. This method sets up peer statuses and the appropriate handlers
-// for each test peer.
-func connectPeers(t *testing.T, host *p2pt.TestP2P, data []*peerData, peerStatus *peers.Status) {
-	const topic = "/eth2/beacon_chain/req/beacon_blocks_by_range/1/ssz"
-
-	for _, d := range data {
-		peer := p2pt.NewTestP2P(t)
-
-		// Copy pointer for callback scope.
-		var datum = d
-
-		peer.SetStreamHandler(topic, func(stream network.Stream) {
-			defer func() {
-				if err := stream.Close(); err != nil {
-					t.Log(err)
-				}
-			}()
-
-			req := &p2ppb.BeaconBlocksByRangeRequest{}
-			if err := peer.Encoding().DecodeWithLength(stream, req); err != nil {
-				t.Error(err)
-			}
-
-			requestedBlocks := makeSequence(req.StartSlot, req.StartSlot+(req.Count*req.Step))
-
-			// Expected failure range
-			if len(sliceutil.IntersectionUint64(datum.failureSlots, requestedBlocks)) > 0 {
-				if _, err := stream.Write([]byte{0x01}); err != nil {
-					t.Error(err)
-				}
-				if _, err := peer.Encoding().EncodeWithLength(stream, "bad"); err != nil {
-					t.Error(err)
-				}
-				return
-			}
-
-			// Determine the correct subset of blocks to return as dictated by the test scenario.
-			blocks := sliceutil.IntersectionUint64(datum.blocks, requestedBlocks)
-
-			ret := make([]*eth.SignedBeaconBlock, 0)
-			for _, slot := range blocks {
-				if (slot-req.StartSlot)%req.Step != 0 {
-					continue
-				}
-				cache.RLock()
-				parentRoot := cache.rootCache[cache.parentSlotCache[slot]]
-				cache.RUnlock()
-				blk := &eth.SignedBeaconBlock{
-					Block: &eth.BeaconBlock{
-						Slot:       slot,
-						ParentRoot: parentRoot[:],
-					},
-				}
-				// If forked peer, give a different parent root.
-				if datum.forkedPeer {
-					newRoot := hashutil.Hash(parentRoot[:])
-					blk.Block.ParentRoot = newRoot[:]
-				}
-				ret = append(ret, blk)
-				currRoot, err := stateutil.BlockRoot(blk.Block)
-				if err != nil {
-					t.Fatal(err)
-				}
-				logrus.Infof("block with slot %d , signing root %#x and parent root %#x", slot, currRoot, parentRoot)
-			}
-
-			if uint64(len(ret)) > req.Count {
-				ret = ret[:req.Count]
-			}
-
-			for i := 0; i < len(ret); i++ {
-				if err := beaconsync.WriteChunk(stream, peer.Encoding(), ret[i]); err != nil {
-					t.Error(err)
-				}
-			}
-		})
-
-		peer.Connect(host)
-
-		peerStatus.Add(new(enr.Record), peer.PeerID(), nil, network.DirOutbound)
-		peerStatus.SetConnectionState(peer.PeerID(), peers.PeerConnected)
-		peerStatus.SetChainState(peer.PeerID(), &p2ppb.Status{
-			ForkDigest:     params.BeaconConfig().GenesisForkVersion,
-			FinalizedRoot:  []byte(fmt.Sprintf("finalized_root %d", datum.finalizedEpoch)),
-			FinalizedEpoch: datum.finalizedEpoch,
-			HeadRoot:       []byte("head_root"),
-			HeadSlot:       datum.headSlot,
-		})
-	}
-}