201 lines
5.6 KiB
Go
201 lines
5.6 KiB
Go
|
// Copyright 2017 The go-ethereum Authors
|
||
|
// This file is part of the go-ethereum library.
|
||
|
//
|
||
|
// The go-ethereum library is free software: you can redistribute it and/or modify
|
||
|
// it under the terms of the GNU Lesser General Public License as published by
|
||
|
// the Free Software Foundation, either version 3 of the License, or
|
||
|
// (at your option) any later version.
|
||
|
//
|
||
|
// The go-ethereum library is distributed in the hope that it will be useful,
|
||
|
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||
|
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||
|
// GNU Lesser General Public License for more details.
|
||
|
//
|
||
|
// You should have received a copy of the GNU Lesser General Public License
|
||
|
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
|
||
|
|
||
|
/*
|
||
|
the p2p/testing package provides a unit test scheme to check simple
|
||
|
protocol message exchanges with one pivot node and a number of dummy peers
|
||
|
The pivot test node runs a node.Service, the dummy peers run a mock node
|
||
|
that can be used to send and receive messages
|
||
|
*/
|
||
|
|
||
|
package testing
|
||
|
|
||
|
import (
|
||
|
"fmt"
|
||
|
"sync"
|
||
|
"testing"
|
||
|
|
||
|
"github.com/ethereum/go-ethereum/log"
|
||
|
"github.com/ethereum/go-ethereum/node"
|
||
|
"github.com/ethereum/go-ethereum/p2p"
|
||
|
"github.com/ethereum/go-ethereum/p2p/discover"
|
||
|
"github.com/ethereum/go-ethereum/p2p/simulations"
|
||
|
"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
|
||
|
"github.com/ethereum/go-ethereum/rpc"
|
||
|
)
|
||
|
|
||
|
// ProtocolTester is the tester environment used for unit testing protocol
|
||
|
// message exchanges. It uses p2p/simulations framework
|
||
|
type ProtocolTester struct {
|
||
|
*ProtocolSession
|
||
|
network *simulations.Network
|
||
|
}
|
||
|
|
||
|
// 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(t *testing.T, id discover.NodeID, n int, run func(*p2p.Peer, p2p.MsgReadWriter) error) *ProtocolTester {
|
||
|
services := adapters.Services{
|
||
|
"test": func(ctx *adapters.ServiceContext) (node.Service, error) {
|
||
|
return &testNode{run}, nil
|
||
|
},
|
||
|
"mock": func(ctx *adapters.ServiceContext) (node.Service, error) {
|
||
|
return newMockNode(), nil
|
||
|
},
|
||
|
}
|
||
|
adapter := adapters.NewSimAdapter(services)
|
||
|
net := simulations.NewNetwork(adapter, &simulations.NetworkConfig{})
|
||
|
if _, err := net.NewNodeWithConfig(&adapters.NodeConfig{
|
||
|
ID: id,
|
||
|
EnableMsgEvents: true,
|
||
|
Services: []string{"test"},
|
||
|
}); err != nil {
|
||
|
panic(err.Error())
|
||
|
}
|
||
|
if err := net.Start(id); err != nil {
|
||
|
panic(err.Error())
|
||
|
}
|
||
|
|
||
|
node := net.GetNode(id).Node.(*adapters.SimNode)
|
||
|
peers := make([]*adapters.NodeConfig, n)
|
||
|
peerIDs := make([]discover.NodeID, n)
|
||
|
for i := 0; i < n; i++ {
|
||
|
peers[i] = adapters.RandomNodeConfig()
|
||
|
peers[i].Services = []string{"mock"}
|
||
|
peerIDs[i] = peers[i].ID
|
||
|
}
|
||
|
events := make(chan *p2p.PeerEvent, 1000)
|
||
|
node.SubscribeEvents(events)
|
||
|
ps := &ProtocolSession{
|
||
|
Server: node.Server(),
|
||
|
IDs: peerIDs,
|
||
|
adapter: adapter,
|
||
|
events: events,
|
||
|
}
|
||
|
self := &ProtocolTester{
|
||
|
ProtocolSession: ps,
|
||
|
network: net,
|
||
|
}
|
||
|
|
||
|
self.Connect(id, peers...)
|
||
|
|
||
|
return self
|
||
|
}
|
||
|
|
||
|
// Stop stops the p2p server
|
||
|
func (self *ProtocolTester) Stop() error {
|
||
|
self.Server.Stop()
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
// Connect brings up the remote peer node and connects it using the
|
||
|
// p2p/simulations network connection with the in memory network adapter
|
||
|
func (self *ProtocolTester) Connect(selfID discover.NodeID, peers ...*adapters.NodeConfig) {
|
||
|
for _, peer := range peers {
|
||
|
log.Trace(fmt.Sprintf("start node %v", peer.ID))
|
||
|
if _, err := self.network.NewNodeWithConfig(peer); err != nil {
|
||
|
panic(fmt.Sprintf("error starting peer %v: %v", peer.ID, err))
|
||
|
}
|
||
|
if err := self.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 := self.network.Connect(selfID, peer.ID); err != nil {
|
||
|
panic(fmt.Sprintf("error connecting to peer %v: %v", peer.ID, err))
|
||
|
}
|
||
|
}
|
||
|
|
||
|
}
|
||
|
|
||
|
// testNode wraps a protocol run function and implements the node.Service
|
||
|
// interface
|
||
|
type testNode struct {
|
||
|
run func(*p2p.Peer, p2p.MsgReadWriter) error
|
||
|
}
|
||
|
|
||
|
func (t *testNode) Protocols() []p2p.Protocol {
|
||
|
return []p2p.Protocol{{
|
||
|
Length: 100,
|
||
|
Run: t.run,
|
||
|
}}
|
||
|
}
|
||
|
|
||
|
func (t *testNode) APIs() []rpc.API {
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
func (t *testNode) Start(server *p2p.Server) error {
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
func (t *testNode) Stop() error {
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
// mockNode is a testNode which doesn't actually run a protocol, instead
|
||
|
// exposing channels so that tests can manually trigger and expect certain
|
||
|
// messages
|
||
|
type mockNode struct {
|
||
|
testNode
|
||
|
|
||
|
trigger chan *Trigger
|
||
|
expect chan *Expect
|
||
|
err chan error
|
||
|
stop chan struct{}
|
||
|
stopOnce sync.Once
|
||
|
}
|
||
|
|
||
|
func newMockNode() *mockNode {
|
||
|
mock := &mockNode{
|
||
|
trigger: make(chan *Trigger),
|
||
|
expect: make(chan *Expect),
|
||
|
err: make(chan error),
|
||
|
stop: make(chan struct{}),
|
||
|
}
|
||
|
mock.testNode.run = mock.Run
|
||
|
return mock
|
||
|
}
|
||
|
|
||
|
// Run is a protocol run function which just loops waiting for tests to
|
||
|
// instruct it to either trigger or expect a message from the peer
|
||
|
func (m *mockNode) Run(peer *p2p.Peer, rw p2p.MsgReadWriter) error {
|
||
|
for {
|
||
|
select {
|
||
|
case trig := <-m.trigger:
|
||
|
m.err <- p2p.Send(rw, trig.Code, trig.Msg)
|
||
|
case exp := <-m.expect:
|
||
|
m.err <- p2p.ExpectMsg(rw, exp.Code, exp.Msg)
|
||
|
case <-m.stop:
|
||
|
return nil
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
func (m *mockNode) Trigger(trig *Trigger) error {
|
||
|
m.trigger <- trig
|
||
|
return <-m.err
|
||
|
}
|
||
|
|
||
|
func (m *mockNode) Expect(exp *Expect) error {
|
||
|
m.expect <- exp
|
||
|
return <-m.err
|
||
|
}
|
||
|
|
||
|
func (m *mockNode) Stop() error {
|
||
|
m.stopOnce.Do(func() { close(m.stop) })
|
||
|
return nil
|
||
|
}
|