1 files changed, 15 insertions, 14 deletions
diff --git a/p2p/testing/protocoltester.go b/p2p/testing/protocoltester.go
index bfa913ff9..1e1752af8 100644
--- a/p2p/testing/protocoltester.go
+++ b/p2p/testing/protocoltester.go
@@ -25,6 +25,7 @@ package testing
 
 import (
 	"bytes"
+	"crypto/ecdsa"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -51,7 +52,7 @@ type ProtocolTester struct {
 // NewProtocolTester constructs a new ProtocolTester
 // it takes as argument the pivot node id, the number of dummy peers and the
 // protocol run function called on a peer connection by the p2p server
-func NewProtocolTester(id enode.ID, nodeCount int, run func(*p2p.Peer, p2p.MsgReadWriter) error) *ProtocolTester {
+func NewProtocolTester(prvkey *ecdsa.PrivateKey, nodeCount int, run func(*p2p.Peer, p2p.MsgReadWriter) error) *ProtocolTester {
 	services := adapters.Services{
 		"test": func(ctx *adapters.ServiceContext) (node.Service, error) {
 			return &testNode{run}, nil
@@ -62,23 +63,30 @@ func NewProtocolTester(id enode.ID, nodeCount int, run func(*p2p.Peer, p2p.MsgRe
 	}
 	adapter := adapters.NewSimAdapter(services)
 	net := simulations.NewNetwork(adapter, &simulations.NetworkConfig{})
-	if _, err := net.NewNodeWithConfig(&adapters.NodeConfig{
-		ID:              id,
+	nodeConfig := &adapters.NodeConfig{
+		PrivateKey:      prvkey,
 		EnableMsgEvents: true,
 		Services:        []string{"test"},
-	}); err != nil {
+	}
+	if _, err := net.NewNodeWithConfig(nodeConfig); err != nil {
 		panic(err.Error())
 	}
-	if err := net.Start(id); err != nil {
+	if err := net.Start(nodeConfig.ID); err != nil {
 		panic(err.Error())
 	}
 
-	node := net.GetNode(id).Node.(*adapters.SimNode)
+	node := net.GetNode(nodeConfig.ID).Node.(*adapters.SimNode)
 	peers := make([]*adapters.NodeConfig, nodeCount)
 	nodes := make([]*enode.Node, nodeCount)
 	for i := 0; i < nodeCount; i++ {
 		peers[i] = adapters.RandomNodeConfig()
 		peers[i].Services = []string{"mock"}
+		if _, err := net.NewNodeWithConfig(peers[i]); err != nil {
+			panic(fmt.Sprintf("error initializing peer %v: %v", peers[i].ID, err))
+		}
+		if err := net.Start(peers[i].ID); err != nil {
+			panic(fmt.Sprintf("error starting peer %v: %v", peers[i].ID, err))
+		}
 		nodes[i] = peers[i].Node()
 	}
 	events := make(chan *p2p.PeerEvent, 1000)
@@ -94,7 +102,7 @@ func NewProtocolTester(id enode.ID, nodeCount int, run func(*p2p.Peer, p2p.MsgRe
 		network:         net,
 	}
 
-	self.Connect(id, peers...)
+	self.Connect(nodeConfig.ID, peers...)
 
 	return self
 }
@@ -108,13 +116,6 @@ func (t *ProtocolTester) Stop() {
 // p2p/simulations network connection with the in memory network adapter
 func (t *ProtocolTester) Connect(selfID enode.ID, peers ...*adapters.NodeConfig) {
 	for _, peer := range peers {
-		log.Trace(fmt.Sprintf("start node %v", peer.ID))
-		if _, err := t.network.NewNodeWithConfig(peer); err != nil {
-			panic(fmt.Sprintf("error starting peer %v: %v", peer.ID, err))
-		}
-		if err := t.network.Start(peer.ID); err != nil {
-			panic(fmt.Sprintf("error starting peer %v: %v", peer.ID, err))
-		}
 		log.Trace(fmt.Sprintf("connect to %v", peer.ID))
 		if err := t.network.Connect(selfID, peer.ID); err != nil {
 			panic(fmt.Sprintf("error connecting to peer %v: %v", peer.ID, err))