diff options
Diffstat (limited to 'core/consensus_test.go')
-rw-r--r-- | core/consensus_test.go | 58 |
1 files changed, 45 insertions, 13 deletions
diff --git a/core/consensus_test.go b/core/consensus_test.go index f677d00..5f61028 100644 --- a/core/consensus_test.go +++ b/core/consensus_test.go @@ -18,6 +18,7 @@ package core import ( + "context" "encoding/json" "testing" "time" @@ -70,7 +71,7 @@ func (n *network) BroadcastAgreementResult( // SendDKGPrivateShare sends PrivateShare to a DKG participant. func (n *network) SendDKGPrivateShare( recv crypto.PublicKey, prvShare *typesDKG.PrivateShare) { - n.conn.send(types.NewNodeID(recv), prvShare) + n.conn.send(n.nID, types.NewNodeID(recv), prvShare) } // BroadcastDKGPrivateShare broadcasts PrivateShare to all DKG participants. @@ -87,8 +88,13 @@ func (n *network) BroadcastDKGPartialSignature( } // ReceiveChan returns a channel to receive messages from DEXON network. -func (n *network) ReceiveChan() <-chan interface{} { - return make(chan interface{}) +func (n *network) ReceiveChan() <-chan types.Msg { + return make(chan types.Msg) +} + +// ReportBadPeer reports that a peer is sending bad message. +func (n *network) ReportBadPeerChan() chan<- interface{} { + return n.conn.s.sink } func (nc *networkConnection) broadcast(from types.NodeID, msg interface{}) { @@ -96,11 +102,11 @@ func (nc *networkConnection) broadcast(from types.NodeID, msg interface{}) { if nID == from { continue } - nc.send(nID, msg) + nc.send(from, nID, msg) } } -func (nc *networkConnection) send(to types.NodeID, msg interface{}) { +func (nc *networkConnection) send(from, to types.NodeID, msg interface{}) { ch, exist := nc.cons[to] if !exist { return @@ -122,12 +128,15 @@ func (nc *networkConnection) send(to types.NodeID, msg interface{}) { } msgCopy = valCopy } - ch <- msgCopy + ch <- types.Msg{ + PeerID: from, + Payload: msgCopy, + } } type networkConnection struct { s *ConsensusTestSuite - cons map[types.NodeID]chan interface{} + cons map[types.NodeID]chan types.Msg } func (nc *networkConnection) newNetwork(nID types.NodeID) *network { @@ -138,14 +147,19 @@ func (nc *networkConnection) newNetwork(nID types.NodeID) *network { } func (nc *networkConnection) setCon(nID types.NodeID, con *Consensus) { - ch := make(chan interface{}, 1000) + ch := make(chan types.Msg, 1000) go func() { for { - msg := <-ch + var msg types.Msg + select { + case msg = <-ch: + case <-nc.s.ctx.Done(): + return + } var err error // Testify package does not support concurrent call. // Use panic() to detact error. - switch val := msg.(type) { + switch val := msg.Payload.(type) { case *types.Block: err = con.preProcessBlock(val) case *types.Vote: @@ -167,13 +181,31 @@ func (nc *networkConnection) setCon(nID types.NodeID, con *Consensus) { type ConsensusTestSuite struct { suite.Suite - conn *networkConnection + ctx context.Context + ctxCancel context.CancelFunc + conn *networkConnection + sink chan interface{} +} + +func (s *ConsensusTestSuite) SetupTest() { + s.ctx, s.ctxCancel = context.WithCancel(context.Background()) + s.sink = make(chan interface{}, 1000) + go func() { + select { + case <-s.sink: + case <-s.ctx.Done(): + return + } + }() +} +func (s *ConsensusTestSuite) TearDownTest() { + s.ctxCancel() } func (s *ConsensusTestSuite) newNetworkConnection() *networkConnection { return &networkConnection{ s: s, - cons: make(map[types.NodeID]chan interface{}), + cons: make(map[types.NodeID]chan types.Msg), } } @@ -384,7 +416,7 @@ func (s *ConsensusTestSuite) TestInitialHeightEventTriggered() { network, prvKey, []*types.Block(nil), - []interface{}(nil), + []types.Msg{}, &common.NullLogger{}, ) s.Require().NoError(err) |